├── .gitignore ├── val_grader ├── __main__.py ├── tests.py └── grader.py ├── data ├── 01.jpg └── 02.jpg ├── run.sh ├── environment.yml ├── job.txt └── project ├── fetchdata.py ├── __init__.py ├── dataloader.py ├── train.py └── models.py /.gitignore: -------------------------------------------------------------------------------- 1 | # IDE files 2 | .idea/* 3 | -------------------------------------------------------------------------------- /val_grader/__main__.py: -------------------------------------------------------------------------------- 1 | from . import grader, tests 2 | grader.run() -------------------------------------------------------------------------------- /data/01.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sftwre/ImageCompressionNN/HEAD/data/01.jpg -------------------------------------------------------------------------------- /data/02.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sftwre/ImageCompressionNN/HEAD/data/02.jpg -------------------------------------------------------------------------------- /run.sh: -------------------------------------------------------------------------------- 1 | python project/train.py --test_path="../flickr30k-images/" --train_path="../flickr30k-images" -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: ml 2 | 3 | dependencies: 4 | - numpy 5 | - pytorch 6 | - pandas 7 | - scipy 8 | - wget 9 | - pillow 10 | -------------------------------------------------------------------------------- /job.txt: -------------------------------------------------------------------------------- 1 | universe = vanilla 2 | Initialdir = /u/isaac3/ml_workspace/ImageCompressionNN/ 3 | Executable = /lusr/bin/bash 4 | Arguments = run.sh 5 | +Group = "GRAD" 6 | +Project = "INSTRUCTIONAL" 7 | +ProjectDescription = "Course project for CS395T Deep Learning Seminar" 8 | Requirements = TARGET.GPUSlot 9 | getenv = True 10 | request_GPUs = 1 11 | +GPUJob = true 12 | Log = condor_ours_sandy.log 13 | Error = condor_ours_sandy.err 14 | Output = condor_ours_sandy.out 15 | Notification = complete 16 | Notify_user = isaac.buitrago.pro@gmail.com 17 | Queue 1 -------------------------------------------------------------------------------- /project/fetchdata.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | import wget 3 | import argparse 4 | import multiprocessing as mp 5 | 6 | def download_kodak(): 7 | 8 | url = "http://r0k.us/graphics/kodak/thumbs/kodim01t.jpg" 9 | pass 10 | 11 | 12 | def fetch_data(urls): 13 | """ 14 | :param urls: list of urls to images 15 | :param save_path: path to place downloaded images 16 | :param num_images: 17 | :return: 18 | """ 19 | # download images 20 | for i, url in enumerate(urls): 21 | print(f"Downloading image {i + 1} from {url}") 22 | wget.download(url, out=save_path) 23 | 24 | 25 | 26 | parser = argparse.ArgumentParser() 27 | 28 | parser.add_argument("--csv", type=str, required=True) 29 | parser.add_argument("--save_path", type=str, required=True) 30 | 31 | args = parser.parse_args() 32 | 33 | save_path = args.save_path 34 | 35 | # create thread pool 36 | num_processes = 10 37 | dataset_size = 1000 38 | 39 | pool = mp.Pool(processes=num_processes) 40 | 41 | # read image urls from csv 42 | df = pd.read_csv(args.csv) 43 | 44 | # compute number of images each process should download 45 | chunksize = dataset_size / num_processes 46 | 47 | pool.map(fetch_data, list(df.TIFF), chunksize=chunksize) 48 | -------------------------------------------------------------------------------- /project/__init__.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import PIL 3 | from models import Encoder, Decoder 4 | from torchvision import transforms 5 | import torch 6 | 7 | TEST_TRANSFORMS_256 = transforms.Compose([ 8 | transforms.ToTensor() 9 | ]) 10 | 11 | encoder = Encoder() 12 | decoder = Decoder() 13 | 14 | def encode(img, bottleneck): 15 | """ 16 | Your code here 17 | img: a 256x256 PIL Image 18 | bottleneck: an integer from {4096,16384,65536} 19 | return: a numpy array less <= bottleneck bytes 20 | """ 21 | 22 | # normalize pixel values 23 | img = img / 255.0 24 | 25 | # conv layer wants 4 dimensions, batch of one image 26 | img = TEST_TRANSFORMS_256(img).unsqueeze(0) 27 | 28 | with torch.no_grad(): 29 | res = encoder.forward(img.cuda()) 30 | 31 | return res 32 | 33 | 34 | def decode(x, bottleneck): 35 | """ 36 | Your code here 37 | x: a numpy array 38 | bottleneck: an integer from {4096,16384,65536} 39 | return a 256x256 PIL Image 40 | """ 41 | 42 | img = torch.from_numpy(x).float().cuda().reshape(1, 64, 32, 32) 43 | 44 | # do inverse of Test Transform 45 | 46 | # need to unormalize an image to visualize it or pass it to the grader 47 | 48 | with torch.no_grad(): 49 | res = decoder.forward(img) 50 | 51 | return res -------------------------------------------------------------------------------- /project/dataloader.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | from torchvision import transforms 4 | from PIL import Image 5 | 6 | 7 | TRAIN_TRANSFORMS_256 = transforms.Compose([ 8 | transforms.RandomResizedCrop(256), 9 | transforms.RandomHorizontalFlip(), 10 | transforms.ToTensor(), 11 | transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) # from [0, 1] to [-1, 1] 12 | ]) 13 | 14 | TEST_TRANSFORMS_256 = transforms.Compose([ 15 | transforms.Resize(256), 16 | transforms.CenterCrop(256), 17 | transforms.ToTensor(), 18 | transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) 19 | ]) 20 | 21 | def is_image_file(img): 22 | return any(img.endswith(ext) for ext in [".png", ".jpg", ".jpeg", ".JPEG"]) 23 | 24 | def load_image(path): 25 | return Image.open(path).convert('RGB') 26 | 27 | def get_training_set(train_path): 28 | return DatasetFromFolder(train_path) 29 | 30 | 31 | class DatasetFromFolder(torch.utils.data.Dataset): 32 | """ 33 | Loads dataset from a given directory 34 | """ 35 | 36 | def __init__(self, train_path): 37 | super(DatasetFromFolder, self).__init__() 38 | self.image_filenames = [os.path.join(train_path, img) for img in os.listdir(train_path) if is_image_file(img)] 39 | 40 | def __getitem__(self, item): 41 | 42 | img = load_image(self.image_filenames[item]) 43 | img = TRAIN_TRANSFORMS_256(img) 44 | 45 | return img 46 | 47 | def __len__(self): 48 | return len(self.image_filenames) 49 | 50 | 51 | -------------------------------------------------------------------------------- /project/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import torch 3 | from dataloader import * 4 | from models import Net 5 | import torch.nn 6 | from torch.utils.data import DataLoader 7 | import os 8 | 9 | # parse command line args 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument('--train_path', help='path to training dataset') 12 | parser.add_argument('--test_path', help='path to testing dataset') 13 | parser.add_argument('--batch_size', type=int, default=10) 14 | parser.add_argument('--num_train', type=int, help='number of training images', default=750) 15 | parser.add_argument('--num_test', type=int, help='number of validation images', default=250) 16 | parser.add_argument('--weights_file', default="network_weights") 17 | parser.add_argument('--lr', default=1e-4, type=float, help='learning rate') 18 | parser.add_argument('--num_epoch', default=80, type=int, help='number of training epochs') 19 | args = parser.parse_args() 20 | 21 | # run model on multiple GPU's 22 | net = torch.nn.DataParallel(Net().cuda()) 23 | 24 | optimizer = torch.optim.Adam(net.parameters(), lr=args.lr, weight_decay=1e-5) 25 | best_loss = torch.tensor(100.0).cuda() 26 | 27 | 28 | def train(dataloader, epoch): 29 | """ 30 | :param dataloader: 31 | :param epoch: 32 | :return: 33 | """ 34 | 35 | for idx, data in enumerate(dataloader): 36 | img = data 37 | img = img.cuda() 38 | 39 | optimizer.zero_grad() 40 | 41 | # encode image 42 | doc = net(img) 43 | 44 | # decode image 45 | doc = net(doc, encode=False) 46 | 47 | loss = torch.mean(torch.abs(img - doc)) 48 | loss.backward() 49 | optimizer.step() 50 | 51 | if idx % 10 == 0: 52 | print(f"Epoch: {epoch} Training Loss: {loss.item()}") 53 | 54 | 55 | def test(dataloader, epoch): 56 | 57 | """ 58 | :param dataloader: 59 | :param epoch: 60 | :return: 61 | """ 62 | 63 | net.eval() 64 | loss = 0 65 | 66 | for idx, data in enumerate(dataloader): 67 | 68 | img = data 69 | 70 | # limit to test set 71 | if idx > args.num_test: 72 | break 73 | 74 | with torch.no_grad(): 75 | img = img.cuda() 76 | doc = net(img, isTest=True) 77 | 78 | # compute validation L1 loss in 0, 1 space 79 | loss += torch.mean(torch.abs(doc - img)) 80 | 81 | loss /= len(dataloader) 82 | print("Epoch: ", epoch, " Testing Loss: ", loss.item()) 83 | 84 | global best_loss 85 | if loss < best_loss: 86 | best_loss = loss.item() 87 | 88 | net.train() 89 | 90 | 91 | def main(): 92 | 93 | print("===> Loading Data") 94 | train_path = os.path.join(os.getcwd(), args.train_path) 95 | train_data = get_training_set(train_path) 96 | print("===> Constructing DataLoader") 97 | 98 | dataloader = DataLoader(dataset=train_data, num_workers=4, batch_size=args.batch_size, shuffle=True) 99 | 100 | print("Training") 101 | # train model 102 | for i in range(args.num_epoch): 103 | train(dataloader, i) 104 | 105 | 106 | # save weights file 107 | torch.save(net.state_dict(), args.weight_file) 108 | 109 | 110 | main() 111 | -------------------------------------------------------------------------------- /val_grader/tests.py: -------------------------------------------------------------------------------- 1 | from .grader import Grader, Case 2 | 3 | import numpy as np 4 | import torch 5 | 6 | def get_data_loader(path_name, batch_size=1): 7 | 8 | from pathlib import Path 9 | from PIL import Image 10 | 11 | path = Path(path_name) 12 | 13 | def _loader(): 14 | for img_path in path.glob('*.jpg'): 15 | img = Image.open(img_path) 16 | yield img 17 | 18 | return _loader 19 | 20 | class PerceptualLoss(torch.nn.Module): 21 | """https://towardsdatascience.com/pytorch-implementation-of-perceptual-losses-for-real-time-style-transfer-8d608e2e9902""" 22 | def __init__(self, vgg): 23 | super().__init__() 24 | self.vgg_features = vgg.features 25 | self.layers = { 26 | '3': "relu1_2", 27 | '8': "relu2_2", 28 | '15': "relu3_3", 29 | '22': "relu4_3" 30 | } 31 | 32 | def forward(self, x): 33 | outputs = dict() 34 | for name, module in self.vgg_features._modules.items(): 35 | x = module(x) 36 | if name in self.layers: 37 | outputs[self.layers[name]] = x.detach() 38 | return outputs 39 | 40 | 41 | class CompressionGrader(Grader): 42 | """Image compression""" 43 | def __init__(self, *args, **kwargs): 44 | bottenecks = [4096,16384,65536] 45 | super().__init__(*args, **kwargs) 46 | 47 | self.encode = self.module.encode 48 | self.decode = self.module.decode 49 | 50 | self.scores = dict() 51 | for bottleneck in bottenecks: 52 | l1, ssim, perceptual = self._get_performance(bottleneck) 53 | self.scores[bottleneck] = (l1, ssim, perceptual) 54 | 55 | print ("[%s B] L1: %.3f, SSIM: %.3f, Perceptual: %.3f"\ 56 | %(bottleneck, np.mean(l1), np.mean(ssim), np.mean(perceptual))) 57 | 58 | 59 | def _get_performance(self, bottleneck): 60 | from itertools import combinations 61 | from skimage.measure import compare_ssim as _compare_ssim 62 | from torchvision.models import vgg 63 | from torchvision.transforms import functional as TF 64 | _vgg16 = vgg.vgg16(pretrained=True) 65 | _perceptual = PerceptualLoss(_vgg16).eval() 66 | _tensor = lambda x: TF.normalize(TF.to_tensor(x),mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])[None] 67 | _numpy = lambda x: np.array(x,dtype=float)/255. 68 | 69 | compare_l1 = lambda a, b: np.abs(_numpy(a) - _numpy(b)).mean() 70 | compare_ssim = lambda a, b: _compare_ssim(_numpy(a), _numpy(b), multichannel=True) 71 | 72 | l1 = [] 73 | ssim = [] 74 | perceptual = [] 75 | 76 | data_loader = get_data_loader('data') 77 | 78 | def compare_perceptual(a, b): 79 | a_features = _perceptual(_tensor(a)) 80 | b_features = _perceptual(_tensor(b)) 81 | loss = 0. 82 | for name in a_features: 83 | loss += float((a_features[name] - b_features[name]).abs().mean()) 84 | 85 | return loss 86 | 87 | for img in data_loader(): 88 | w, h = img.size 89 | 90 | z = self.encode(img, bottleneck) 91 | 92 | assert z.nbytes <= bottleneck, "Latent vector exceeds bottleneck" 93 | 94 | img_rec = self.decode(z, bottleneck) 95 | 96 | # from PIL import Image 97 | # size = {4096:36,16384:73,65536:147}.get(bottleneck) 98 | # img_low = img.resize((size,size),Image.ANTIALIAS) 99 | # img_rec = img_low.resize((256,256),Image.ANTIALIAS) 100 | 101 | assert img_rec.size == img.size, "Decoded image has wrong resolution" 102 | 103 | rec_l1 = compare_l1(img, img_rec) 104 | rec_ssim = compare_ssim(img, img_rec) 105 | rec_perceptual = compare_perceptual(img, img_rec) 106 | 107 | l1.append(rec_l1) 108 | ssim.append(rec_ssim) 109 | perceptual.append(rec_perceptual) 110 | 111 | return l1, ssim, perceptual 112 | 113 | 114 | @Case(score=10) 115 | def test_low_l1(self, low=0.023, high=0.035): 116 | """4096B: L1 distance""" 117 | return np.clip(high-np.mean(self.scores[4096][0]), 0, high-low) / (high-low) 118 | 119 | @Case(score=10) 120 | def test_med_l1(self, low=0.011, high=0.023): 121 | """16384B: L1 distance""" 122 | return np.clip(high-np.mean(self.scores[16384][0]), 0, high-low) / (high-low) 123 | 124 | @Case(score=10) 125 | def test_high_l1(self, low=0.005, high=0.011): 126 | """65536B: L1 distance""" 127 | return np.clip(high-np.mean(self.scores[65536][0]), 0, high-low) / (high-low) 128 | 129 | @Case(score=10) 130 | def test_low_ssim(self, low=0.788, high=0.892): 131 | """4096B: SSIM""" 132 | return np.clip(np.mean(self.scores[4096][1])-low, 0, high-low) / (high-low) 133 | 134 | @Case(score=10) 135 | def test_med_ssim(self, low=0.892, high=0.971): 136 | """16384B: SSIM""" 137 | return np.clip(np.mean(self.scores[16384][1])-low, 0, high-low) / (high-low) 138 | 139 | @Case(score=10) 140 | def test_high_ssim(self, low=0.971, high=0.99): 141 | """65536B: SSIM""" 142 | return np.clip(np.mean(self.scores[65536][1])-low, 0, high-low) / (high-low) 143 | 144 | @Case(score=10) 145 | def test_low_perceptual(self, low=2.075, high=2.601): 146 | """4096B: perceptual""" 147 | return np.clip(high-np.mean(self.scores[4096][2]), 0, high-low) / (high-low) 148 | 149 | @Case(score=10) 150 | def test_med_perceptual(self, low=1.115, high=2.075): 151 | """16384B: perceptual""" 152 | return np.clip(high-np.mean(self.scores[16384][2]), 0, high-low) / (high-low) 153 | 154 | @Case(score=10) 155 | def test_high_perceptual(self, low=0.9, high=1.115): 156 | """65536B: perceptual""" 157 | return np.clip(high-np.mean(self.scores[65536][2]), 0, high-low) / (high-low) 158 | -------------------------------------------------------------------------------- /project/models.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import numpy as np 3 | import torch 4 | import math 5 | 6 | 7 | class Net(nn.Module): 8 | 9 | def __init__(self): 10 | super(Net, self).__init__() 11 | 12 | self.encoder = Encoder() 13 | self.decoder = Decoder() 14 | 15 | def forward(self, img, encode=True): 16 | 17 | if encode: 18 | return self.encoder(img) 19 | else: 20 | return self.decoder(img) 21 | 22 | 23 | class Encoder(nn.Module): 24 | 25 | 26 | def __init__(self): 27 | 28 | super(Encoder, self).__init__() 29 | 30 | self.scales = 6 31 | self.alignment_scale = (32, 32) 32 | 33 | # decomposition layer 34 | self.decompLayer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1), 35 | nn.LeakyReLU(0.2), 36 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), 37 | nn.LeakyReLU(0.2), 38 | nn.BatchNorm2d(64)) 39 | 40 | # interscale alignment layer 41 | self.downsampleLayers = { 42 | 8: nn.Conv2d(64, 64, kernel_size=3, stride=8, padding=1), 43 | 4: nn.Conv2d(64, 64, kernel_size=3, stride=4, padding=1), 44 | 2: nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1), 45 | 1: nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1) 46 | } 47 | 48 | self.upsampleLayers = { 49 | 2: nn.ConvTranspose2d(64, 64, kernel_size=3, stride=2, padding=1), 50 | 4: nn.ConvTranspose2d(64, 64, kernel_size=3, stride=4, padding=1) 51 | } 52 | 53 | # output layer 54 | self.outputLayer = nn.Sequential(nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), 55 | nn.LeakyReLU(0.2), 56 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), 57 | nn.LeakyReLU(0.2)) 58 | 59 | self.scale_factor = 0.5 60 | self.coef_maps = list() 61 | 62 | 63 | def decompose(self, xm, size): 64 | """ 65 | Performs pyramidal decomposition by extracting 66 | coefficients from the input scale and computing next scale. 67 | 68 | :param xm: Tensor for image 69 | :param size: tuple of height and width desired in interpolation 70 | :return: Tensor for coefficient and Tensor for downsampled image 71 | """ 72 | 73 | # downsample to next scale 74 | xm1 = nn.functional.interpolate(xm, mode='bilinear', size=size) 75 | xm = self.decompLayer(xm) 76 | 77 | # return coefficiant and downsampled image 78 | return xm, xm1 79 | 80 | 81 | def align(self): 82 | """ 83 | Performs interscale alignment of features in the 84 | coef_map. Computes difference between size of coef tensor 85 | and alignment_scale then passes coef through appropriate 86 | conv layer. coef_map must contain a tensor for each scale. 87 | 88 | :returns Tensor for sum of coefficients 89 | """ 90 | # len(self.coef_maps) == self.scales 91 | print(len(self.coef_maps)) 92 | 93 | # sum of coefficient tensors 94 | y = torch.zeros(size=(32, 32, 64), dtype=torch.float32) 95 | 96 | for coef in self.coef_maps: 97 | 98 | # dimensions of coef tensor and desired alignment 99 | align_scale = self.alignment_scale 100 | coef_scale = tuple(coef.size()) 101 | 102 | # determine which conv to pass img through 103 | if coef_scale > self.alignment_scale: 104 | conv = self.downsampleLayers[int(coef_scale[0] / align_scale[0])] 105 | else: 106 | 107 | print(coef_scale) 108 | conv = self.upsampleLayers[int(align_scale[0] / coef_scale[0])] 109 | 110 | # align coefficients 111 | y += conv(coef) 112 | 113 | return y 114 | 115 | 116 | def forward(self, x): 117 | """ 118 | :param x Tensor for Image that will be compressed 119 | :returns Tensor for compressed image 120 | """ 121 | 122 | xm = x 123 | dimensions = np.array([256.0, 256.0]) 124 | 125 | # perform pyramidal decomposition 126 | for scale in range(self.scales): 127 | x, xm = self.decompose(xm, tuple(map(lambda x: int(x), dimensions))) 128 | self.coef_maps.append(x) 129 | dimensions *= self.scale_factor 130 | print(dimensions) 131 | 132 | 133 | # perform interscale alignment 134 | y = self.align() 135 | 136 | # convolve aligned features 137 | y = self.outputLayer(y) 138 | 139 | # compressed image 140 | return y 141 | 142 | 143 | class Quantization(nn.Module): 144 | 145 | B = 6 146 | 147 | def __init__(self): 148 | super(Quantization, self).__init__() 149 | 150 | def forward(self, y): 151 | return (1 / pow(2, self.B - 1)) * math.ceil(pow(2, self.B - 1) * y) 152 | 153 | 154 | class Decoder(nn.Module): 155 | 156 | def __init__(self): 157 | super(Decoder, self).__init__() 158 | 159 | self.layer1 = nn.Sequential(nn.ConvTranspose2d(64, 64, kernel_size=3, stride=8, padding=1), nn.LeakyReLU(0.2)) 160 | self.layer2 = nn.Sequential(nn.ConvTranspose2d(64, 64, kernel_size=3, stride=8, padding=1), nn.LeakyReLU(0.2)) 161 | self.layer3 = nn.Sequential(nn.ConvTranspose2d(64, 64, kernel_size=3, stride=8, padding=1), nn.LeakyReLU(0.2)) 162 | self.layer4 = nn.Sequential(nn.ConvTranspose2d(64, 64, kernel_size=3, stride=8, padding=1), nn.LeakyReLU(0.2)) 163 | 164 | 165 | def forward(self, img): 166 | """ 167 | 168 | :param img: 32x32 compressed image 169 | :return: 256X256 reconstructed image 170 | """ 171 | 172 | 173 | img = self.layer1(img) 174 | img = nn.functional.interpolate(img, mode="bilinear", size=(64, 64, 64)) 175 | img = self.layer2(img) 176 | img = nn.functional.interpolate(img, mode="bilinear", size=(128, 128, 64)) 177 | img = self.layer3(img) 178 | img = nn.functional.interpolate(img, mode="bilinear", size=(256, 256, 64)) 179 | img = self.layer4(img) 180 | 181 | return img 182 | -------------------------------------------------------------------------------- /val_grader/grader.py: -------------------------------------------------------------------------------- 1 | class CheckFailed(Exception): 2 | def __init__(self, why): 3 | self.why = why 4 | 5 | def __str__(self): 6 | return self.why 7 | 8 | 9 | class ContextManager: 10 | def __init__(self, on, off): 11 | self.on = on 12 | self.off = off 13 | 14 | def __enter__(self): 15 | self.on() 16 | 17 | def __exit__(self, exc_type, exc_value, traceback): 18 | self.off() 19 | 20 | 21 | def list_all_kwargs(**kwargs): 22 | all_args = [{}] 23 | for k, v in kwargs.items(): 24 | new_args = [] 25 | for i in v: 26 | new_args.extend([dict({k: i}, **a) for a in all_args]) 27 | all_args = new_args 28 | return all_args 29 | 30 | 31 | # Use @Case(score, extra_credit) as a decorator for member functions of a Grader 32 | # this will make them test cases 33 | # A test case can return a value between 0 and 1 as a score. If the test fails it should raise an assertion. 34 | # The test case may optionally return a tuple (score, message). 35 | 36 | 37 | def case(func, kwargs={}, score=1, extra_credit=False): 38 | def wrapper(self): 39 | msg = 'passed' 40 | n_passed, total = 0.0, 0.0 41 | for a in list_all_kwargs(**kwargs): 42 | try: 43 | v = func(self, **a) 44 | if v is None: 45 | v = 1 46 | elif isinstance(v, tuple): 47 | v, msg = v 48 | else: 49 | assert isinstance(v, float), "case returned %s which is not a float!" % repr(v) 50 | n_passed += v 51 | except AssertionError as e: 52 | msg = str(e) 53 | except CheckFailed as e: 54 | msg = str(e) 55 | except NotImplementedError as e: 56 | msg = 'Function not implemented %s' % e 57 | except Exception as e: 58 | msg = 'Crash "%s"' % e 59 | total += 1 60 | return int(n_passed * score / total + 0.5), msg 61 | 62 | wrapper.score = score 63 | wrapper.extra_credit = extra_credit 64 | wrapper.__doc__ = func.__doc__ 65 | return wrapper 66 | 67 | 68 | class Case(object): 69 | def __init__(self, score=1, extra_credit=False): 70 | self.score = score 71 | self.extra_credit = extra_credit 72 | 73 | def __call__(self, func): 74 | return case(func, score=self.score, extra_credit=self.extra_credit) 75 | 76 | 77 | class MultiCase(object): 78 | def __init__(self, score=1, extra_credit=False, **kwargs): 79 | self.score = score 80 | self.extra_credit = extra_credit 81 | self.kwargs = kwargs 82 | 83 | def __call__(self, func): 84 | return case(func, kwargs=self.kwargs, score=self.score, extra_credit=self.extra_credit) 85 | 86 | 87 | class Grader: 88 | def __init__(self, module, verbose=False): 89 | self.module = module 90 | self.verbose = verbose 91 | 92 | @classmethod 93 | def has_cases(cls): 94 | import inspect 95 | for n, f in inspect.getmembers(cls): 96 | if hasattr(f, 'score'): 97 | return True 98 | return False 99 | 100 | @classmethod 101 | def total_score(cls): 102 | import inspect 103 | r = 0 104 | for n, f in inspect.getmembers(cls): 105 | if hasattr(f, 'score'): 106 | r += f.score 107 | return r 108 | 109 | def run(self): 110 | import inspect 111 | score, total_score = 0, 0 112 | if self.verbose: 113 | print(' * %-50s' % self.__doc__) 114 | for n, f in inspect.getmembers(self): 115 | if hasattr(f, 'score'): 116 | s, msg = f() 117 | score += s 118 | if self.verbose: 119 | print(' - %-50s [ %s ]' % (f.__doc__, msg)) 120 | if not f.extra_credit: 121 | total_score += f.score 122 | 123 | return score, total_score 124 | 125 | 126 | def grade(G, assignment_module, verbose=False): 127 | try: 128 | grader = G(assignment_module, verbose) 129 | except NotImplementedError as e: 130 | if verbose: 131 | print(' - Function not implemented: %s' % e) 132 | return 0, G.total_score() 133 | except Exception as e: 134 | if verbose: 135 | print(' - Your program crashed "%s"' % e) 136 | return 0, G.total_score() 137 | 138 | return grader.run() 139 | 140 | 141 | def grade_all(assignment_module, verbose=False): 142 | score, total_score = 0, 0 143 | for G in Grader.__subclasses__(): 144 | if G.has_cases(): 145 | s, ts = grade(G, assignment_module, verbose) 146 | 147 | if verbose: 148 | print(' -------------------------------------------------- [ %3d / %3d ]' % (s, ts)) 149 | print() 150 | else: 151 | print(' * %-50s [ %3d / %3d ]' % (G.__doc__, s, ts)) 152 | total_score += ts 153 | score += s 154 | 155 | print() 156 | print('total score %3d / %3d' % (score, total_score)) 157 | 158 | 159 | def load_assignment(name): 160 | import atexit 161 | from glob import glob 162 | import importlib 163 | from os import path 164 | from shutil import rmtree 165 | import sys 166 | import tempfile 167 | import zipfile 168 | 169 | if path.isdir(name): 170 | return importlib.import_module(name) 171 | 172 | with zipfile.ZipFile(name) as f: 173 | tmp_dir = tempfile.mkdtemp() 174 | atexit.register(lambda: rmtree(tmp_dir)) 175 | 176 | f.extractall(tmp_dir) 177 | module_names = glob(path.join(tmp_dir, '*')) 178 | assert len(module_names) == 1, 'Malformed zip file, expecting exactly one top-level folder, got %d' % \ 179 | len(module_names) 180 | sys.path.insert(0, tmp_dir) 181 | module = path.basename(module_names[0]) 182 | return importlib.import_module(module) 183 | 184 | 185 | def run(): 186 | import argparse 187 | 188 | parser = argparse.ArgumentParser('Grade your assignment') 189 | parser.add_argument('assignment', default='homework') 190 | parser.add_argument('-v', '--verbose', action='store_true') 191 | args = parser.parse_args() 192 | 193 | print('Loading assignment') 194 | assignment = load_assignment(args.assignment) 195 | 196 | print('Loading grader') 197 | grade_all(assignment, args.verbose) --------------------------------------------------------------------------------