Update app.py

app.py

@@ -4,6 +4,10 @@ from PIL import Image
 import pytesseract
 import yolov5
 
+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')
 
@@ -18,30 +22,61 @@ def license_plate_detect(img):
     results = model(img, size=640)
 
     # parse results
-    if len(results.pred):
-        predictions = results.pred[0]
+    predictions = results.pred[0]
+    if len(predictions):
         boxes = predictions[:, :4] # x1, y1, x2, y2
         return boxes
 
 
 def read_license_number(img):
     boxes = license_plate_detect(img)
-    if len(boxes[0]):
+    if boxes:
         image = Image.fromarray(img)
         return [pytesseract.image_to_string(
                     image.crop(bbox.tolist()))
                for bbox in boxes]
 
 
+def zero_shot_classification(image, labels):
+    inputs = processor(text=labels,
+                       images=image,
+                       return_tensors="pt",
+                       padding=True)
+    outputs = model(**inputs)
+    logits_per_image = outputs.logits_per_image # this is the image-text similarity score
+    return logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities
+
+installed_list = []
+# image = Image.open(requests.get(url, stream=True).raw)
+def check_solarplant_installed_by_license(license_number_list):
+    if len(installed_list):
+        return [license_number in installed_list
+                for license_number in license_number_list]
+
+def check_solarplant_installed_by_image(image, output_label=False):
+    zero_shot_class_labels = ["bus with solar panel grids",
+                              "bus without solar panel grids"]
+    probs = zero_shot_classification(image, zero_shot_class_labels)
+    if output_label:
+        return zero_shot_class_labels[probs.argmax().item()]
+    return probs.argmax().item() == 0
+
+def check_solarplant_broken(image):
+    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]
+
+
 def greet(img):
     boxes = license_plate_detect(img)
-    image = Image.fromarray(img)
-    r = 'greet'
-    if len(boxes[0]):
-        r = [pytesseract.image_to_string(
-                    image.crop(bbox.tolist()))
-               for bbox in boxes]
-    return "Hello " + str(r) + "!!"
+    if boxes:
+        return (seg,
+                "車牌： " + '; '.join(lns) + "\n\n" \
+                + "類型： "+ check_solarplant_installed_by_image(img, True) + "\n\n" \
+                + "狀態：" + check_solarplant_broken(img))
+    return (img, "空地。。。")
 
 
 iface = gr.Interface(fn=greet, inputs="image", outputs="text")