npc0 commited on
Commit
fd69589
1 Parent(s): 41af20c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -10
app.py CHANGED
@@ -4,6 +4,10 @@ from PIL import Image
4
  import pytesseract
5
  import yolov5
6
 
 
 
 
 
7
  # load model
8
  model = yolov5.load('keremberke/yolov5m-license-plate')
9
 
@@ -18,30 +22,61 @@ def license_plate_detect(img):
18
  results = model(img, size=640)
19
 
20
  # parse results
21
- if len(results.pred):
22
- predictions = results.pred[0]
23
  boxes = predictions[:, :4] # x1, y1, x2, y2
24
  return boxes
25
 
26
 
27
  def read_license_number(img):
28
  boxes = license_plate_detect(img)
29
- if len(boxes[0]):
30
  image = Image.fromarray(img)
31
  return [pytesseract.image_to_string(
32
  image.crop(bbox.tolist()))
33
  for bbox in boxes]
34
 
35
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
36
  def greet(img):
37
  boxes = license_plate_detect(img)
38
- image = Image.fromarray(img)
39
- r = 'greet'
40
- if len(boxes[0]):
41
- r = [pytesseract.image_to_string(
42
- image.crop(bbox.tolist()))
43
- for bbox in boxes]
44
- return "Hello " + str(r) + "!!"
45
 
46
 
47
  iface = gr.Interface(fn=greet, inputs="image", outputs="text")
 
4
  import pytesseract
5
  import yolov5
6
 
7
+ from transformers import CLIPProcessor, CLIPModel
8
+ model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
9
+ processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
10
+
11
  # load model
12
  model = yolov5.load('keremberke/yolov5m-license-plate')
13
 
 
22
  results = model(img, size=640)
23
 
24
  # parse results
25
+ predictions = results.pred[0]
26
+ if len(predictions):
27
  boxes = predictions[:, :4] # x1, y1, x2, y2
28
  return boxes
29
 
30
 
31
  def read_license_number(img):
32
  boxes = license_plate_detect(img)
33
+ if boxes:
34
  image = Image.fromarray(img)
35
  return [pytesseract.image_to_string(
36
  image.crop(bbox.tolist()))
37
  for bbox in boxes]
38
 
39
 
40
+ def zero_shot_classification(image, labels):
41
+ inputs = processor(text=labels,
42
+ images=image,
43
+ return_tensors="pt",
44
+ padding=True)
45
+ outputs = model(**inputs)
46
+ logits_per_image = outputs.logits_per_image # this is the image-text similarity score
47
+ return logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities
48
+
49
+ installed_list = []
50
+ # image = Image.open(requests.get(url, stream=True).raw)
51
+ def check_solarplant_installed_by_license(license_number_list):
52
+ if len(installed_list):
53
+ return [license_number in installed_list
54
+ for license_number in license_number_list]
55
+
56
+ def check_solarplant_installed_by_image(image, output_label=False):
57
+ zero_shot_class_labels = ["bus with solar panel grids",
58
+ "bus without solar panel grids"]
59
+ probs = zero_shot_classification(image, zero_shot_class_labels)
60
+ if output_label:
61
+ return zero_shot_class_labels[probs.argmax().item()]
62
+ return probs.argmax().item() == 0
63
+
64
+ def check_solarplant_broken(image):
65
+ zero_shot_class_labels = ["white broken solar panel",
66
+ "normal black solar panel grids"]
67
+ probs = zero_shot_classification(image, zero_shot_class_labels)
68
+ idx = probs.argmax().item()
69
+ return zero_shot_class_labels[idx][1-idx]
70
+
71
+
72
  def greet(img):
73
  boxes = license_plate_detect(img)
74
+ if boxes:
75
+ return (seg,
76
+ "車牌: " + '; '.join(lns) + "\n\n" \
77
+ + "類型: "+ check_solarplant_installed_by_image(img, True) + "\n\n" \
78
+ + "狀態:" + check_solarplant_broken(img))
79
+ return (img, "空地。。。")
 
80
 
81
 
82
  iface = gr.Interface(fn=greet, inputs="image", outputs="text")