Update app.py

app.py

@@ -7,15 +7,15 @@ from transformers import CLIPProcessor, CLIPModel
 model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
 processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
-# # load model
-# model = yolov5.load('keremberke/yolov5m-license-plate')
+# load model
+model = yolov5.load('keremberke/yolov5m-license-plate')
 
-# # set model parameters
-# model.conf = 0.5  # NMS confidence threshold
-# model.iou = 0.25  # NMS IoU threshold
-# model.agnostic = False  # NMS class-agnostic
-# model.multi_label = False  # NMS multiple labels per box
-# model.max_det = 1000  # maximum number of detections per image
+# set model parameters
+model.conf = 0.5  # NMS confidence threshold
+model.iou = 0.25  # NMS IoU threshold
+model.agnostic = False  # NMS class-agnostic
+model.multi_label = False  # NMS multiple labels per box
+model.max_det = 1000  # maximum number of detections per image
 
 def license_plate_detect(img):
     results = model(img, size=640)
@@ -62,23 +62,26 @@ def check_solarplant_installed_by_image(image, output_label=False):
         return zero_shot_class_labels[probs.argmax().item()]
     return probs.argmax().item() == 0
 
-def check_solarplant_broken(image):
+def check_solarplant_broken(image, output_label=False):
     zero_shot_class_labels = ["white broken solar panel",
                               "normal black solar panel grids"]
     probs = zero_shot_classification(image, zero_shot_class_labels)
     idx = probs.argmax().item()
-    return zero_shot_class_labels[idx][1-idx]
+    if output_label:
+        return zero_shot_class_labels[idx].split(" ")[1-idx]
+    return idx == 0
 
 
 def greet(img):
     print(type(img))
-    # lns = read_license_number(img)
-    lns = ['1','2','3']
+    lns = read_license_number(img)
     if len(lns):
+        planttype = check_solarplant_installed_by_image(img, True)
         # return (seg,
         return ("車牌： " + '; '.join(lns) + "\n\n" \
-                + "類型： "+ check_solarplant_installed_by_image(img, True) + "\n\n" \
-                + "狀態：" + check_solarplant_broken(img))
+                + "類型： "+ planttype + "\n\n" \
+                + "狀態：" + (check_solarplant_broken(img, True)
+                           if 'with' in planttype else 'normal'))
     return (img, "空地。。。")