Spaces:
Running
Running
| import numpy as np | |
| import cv2 | |
| import torch | |
| import glob as glob | |
| import os | |
| import matplotlib.pyplot as plt | |
| from models.create_fasterrcnn_model import create_model | |
| from utils.annotations import CNNpostAnnotations | |
| #from utils.annotations import inference_annotations | |
| from utils.general import set_infer_dir | |
| from utils.transforms import infer_transforms | |
| import numpy as np | |
| from skimage import transform | |
| import os | |
| from keras.models import Model | |
| from keras.optimizers import Adam | |
| from keras.applications.vgg16 import VGG16, preprocess_input | |
| from keras.layers import Dense, Dropout, Flatten | |
| import numpy as np | |
| conv_base = VGG16(include_top=False, | |
| weights='imagenet', | |
| input_shape=(200,200,3)) | |
| if 2 > 0: | |
| for layer in conv_base.layers[:-2]: | |
| layer.trainable = False | |
| else: | |
| for layer in conv_base.layers: | |
| layer.trainable = False | |
| top_model = conv_base.output | |
| top_model = Flatten(name="flatten")(top_model) | |
| top_model = Dense(4096, activation='relu')(top_model) | |
| top_model = Dense(1048, activation='relu')(top_model) | |
| top_model = Dense(256, activation='relu')(top_model) | |
| top_model = Dense(128, activation='relu')(top_model) | |
| top_model = Dense(64, activation='relu')(top_model) | |
| top_model = Dropout(0.2)(top_model) | |
| output_layer = Dense(5, activation='softmax')(top_model) | |
| CNN = Model(inputs=conv_base.input, outputs=output_layer) | |
| CNN.load_weights("CNN.hdf5") | |
| def main(weightUrl, input): | |
| np.random.seed(42) | |
| NUM_CLASSES = 2 | |
| CLASSES = ['__background__', 'Cell'] | |
| DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') | |
| OUT_DIR = set_infer_dir() | |
| checkpoint = torch.load(weightUrl, map_location=DEVICE) | |
| data_configs = True | |
| NUM_CLASSES = checkpoint['config']['NC'] | |
| CLASSES = checkpoint['config']['CLASSES'] | |
| build_model = create_model[checkpoint['model_name']] | |
| model = build_model(num_classes=NUM_CLASSES, coco_model=False) | |
| model.load_state_dict(checkpoint['model_state_dict']) | |
| model.to(DEVICE).eval() | |
| COLORS = np.random.uniform(0, 255, size=(len(CLASSES), 3)) | |
| image = input | |
| orig_image = image.copy() | |
| image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB) | |
| image = infer_transforms(image) | |
| image = torch.unsqueeze(image, 0) | |
| outputs = model(image.to(DEVICE)) | |
| # Load all detection to CPU for further operations. | |
| outputs = [{k: v.to('cpu') for k, v in t.items()} for t in outputs] | |
| print(outputs) | |
| # Carry further only if there are detected boxes. | |
| if len(outputs[0]['boxes']) != 0: | |
| # orig_image = inference_annotations( | |
| # outputs, 0.3, CLASSES, | |
| # (255, 255, 255), orig_image | |
| # ) | |
| orig_image, cellImgs = CNNpostAnnotations( | |
| outputs, 0.3, CLASSES, | |
| (255, 255, 255), orig_image, CNN | |
| ) | |
| return orig_image, cellImgs | |
| cv2.destroyAllWindows() | |