File size: 3,189 Bytes
4e1708d b159813 4e1708d b159813 4e1708d b159813 4e1708d b159813 4e1708d f2d04f5 4e1708d b159813 4e1708d b159813 4e1708d b159813 4e1708d b159813 4e1708d b159813 4e1708d | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | import os
import gradio as gr
from fastai.vision.all import *
from icevision.all import *
from icevision.models.checkpoint import *
print("Loading images")
for root, dirs, files in os.walk(r"sample_images/"):
for filename in files:
print(filename)
print("Loading classifier")
classifier = load_learner("models/learner.pkl")
classifier_labels = classifier.dls.vocab
print("Loading detector")
checkpoint_path = "models/model_checkpoint.pth"
checkpoint_and_model = model_from_checkpoint(checkpoint_path)
model = checkpoint_and_model["model"]
model_type = checkpoint_and_model["model_type"]
class_map = checkpoint_and_model["class_map"]
img_size = checkpoint_and_model["img_size"]
valid_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(img_size), tfms.A.Normalize()])
def draw_eyes(img):
pred_dict = model_type.end2end_detect(
img,
valid_tfms,
model,
class_map=class_map,
detection_threshold=0.5,
display_label=False,
color_map={"in": "#FF4040", "out": "#FFC71E"},
)
# Draw bbox with cv
for i, bbox in enumerate(pred_dict["detection"]["bboxes"]):
x, y, w, h = pred_dict["detection"]["bboxes"][i].xywh
xmin, ymin, xmax, ymax = pred_dict["detection"]["bboxes"][i].xyxy
center = (int((xmin + xmax) / 2), int((ymin + ymax) / 2))
if pred_dict["detection"]["labels"][i] == "out":
color_value = (255, 0, 0)
else:
color_value = (8, 39, 245)
image = cv2.rectangle(np.array(img), (x, y), (x + w, y + h), color_value, 2)
image = cv2.circle(image, center, 5, color_value, -1)
image = cv2.putText(
image,
f"w:{w} h:{h}",
(x, y - 10),
cv2.FONT_HERSHEY_SIMPLEX,
1,
color_value,
2,
cv2.LINE_AA,
)
img = Image.fromarray(image)
return img
def predict(img):
img = PILImage.create(img)
pred, pred_idx, probs = classifier.predict(img)
img = draw_eyes(img)
return {
classifier_labels[i]: float(probs[i]) for i in range(len(classifier_labels))
}, img
title = "NSWX Electrode Classifier"
description = "Upload an image of a bare electrode or select from the examples below"
interpretation = "default"
examples = ["sample_images/" + file for file in files]
article = "<p style='text-align: center'><a href='https://dicksonneoh.com/' target='_blank'>Blog post</a></p>"
enable_queue = False
gr.Interface(
fn=predict,
inputs=gr.inputs.Image( label="Input image"),
outputs=[
gr.outputs.Label(num_top_classes=5, label="Electrode Class"),
gr.outputs.Image(type="pil", label="WE Dimensions"),
],
title=title,
description=description,
article=article,
examples=examples,
interpretation=interpretation,
enable_queue=enable_queue,
allow_flagging="manual",
flagging_options=["This should be OK", "This should be KIV_COL", "This should be KIV_CMT", "This should be NG_DIM", "This should be NG_MSA"],
theme="grass",
css = ".output-image, .input-image, .image-preview {height: 600px !important} ",
).launch(server_name="0.0.0.0")
|