Spaces:
Runtime error
Runtime error
File size: 3,018 Bytes
6b2d891 9a30e62 2e4ea99 adca53b 1094aba 2e4ea99 458c5d4 6b2d891 2e4ea99 fa10257 a64c8d1 2fd0fb2 958d113 2e4ea99 d01979a 2e4ea99 b8316ce d54d7ce 2e4ea99 d01979a cacc64e 2e4ea99 dbe5043 2e4ea99 1094aba 146be9b 2e4ea99 6f792de 2e4ea99 d01979a b9f84d6 d01979a 2e4ea99 612a23d 2e4ea99 d01979a 2e4ea99 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import os
import gradio as gr
import torch
from torchvision import models, transforms
from PIL import Image
# -- get torch and cuda version
#TORCH_VERSION = ".".join(torch.__version__.split(".")[:2])
#CUDA_VERSION = torch.__version__.split("+")[-1]
# -- install pre-build detectron2
os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
os.system('pip install pyyaml==5.1')
os.system('pip install opencv')
import detectron2
from detectron2.utils.logger import setup_logger # ????
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
# ????
setup_logger()
# -- load rcnn model
cfg = get_cfg()
# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold for this model
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
'''
os.system(wget http://images.cocodataset.org/val2017/000000439715.jpg -q -O input.jpg)
im = cv2.imread("./input.jpg")
cv2_imshow(im)
outputs = predictor(im)
print(outputs["instances"].pred_classes)
print(outputs["instances"].pred_boxes)
'''
# -- load Mask R-CNN model for segmentation
DesignModernityModel = torch.load("DesignModernityModel.pt")
#INPUT_FEATURES = DesignModernityModel.fc.in_features
#linear = nn.linear(INPUT_FEATURES, 5)
DesignModernityModel.eval() # set state of the model to inference
LABELS = ['2000-2003', '2006-2008', '2009-2011', '2012-2014', '2015-2018']
n_labels = len(LABELS)
MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]
carTransforms = transforms.Compose([transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize(mean=MEAN, std=STD)])
def classifyCar(im):
im = Image.fromarray(im.astype('uint8'), 'RGB')
im = carTransforms(im).unsqueeze(0) # transform and add batch dimension
with torch.no_grad():
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
scores = torch.nn.functional.softmax(DesignModernityModel(im)[0])
return Image.fromarray(np.uint8(out.get_image())).convert('RGB'), {LABELS[i]: float(scores[i]) for i in range(n_labels)}
#examples = [[example_img.jpg], [example_img2.jpg]] # must be uploaded in repo
# create interface for model
interface = gr.Interface(classifyCar, inputs='image', outputs=['image','label'], cache_examples=False, title='VW Up or Fiat 500')
interface.launch()
|