Soham Chandratre commited on
Commit
59c748e
·
1 Parent(s): ae30f0b

minor changes

Browse files
model/__pycache__/pothole_model.cpython-311.pyc CHANGED
Binary files a/model/__pycache__/pothole_model.cpython-311.pyc and b/model/__pycache__/pothole_model.cpython-311.pyc differ
 
model/pothole_model.py CHANGED
@@ -1,68 +1,69 @@
1
- # from ultralyticsplus import YOLO, render_result
2
- # from PIL import Image
3
- # import numpy as np
 
4
 
5
- # def load_model(image):
6
- # # image_bytes = image.content
7
- # model = YOLO('keremberke/yolov8n-pothole-segmentation')
8
- # model.overrides['conf'] = 0.25
9
- # model.overrides['iou'] = 0.45
10
- # model.overrides['agnostic_nms'] = False
11
- # model.overrides['max_det'] = 1000
12
 
13
- # # Load image using PIL
14
- # image = Image.open((image))
15
- # image_array = np.array(image)
16
- # # pil_image = pil_image.convert("RGB") # Ensure image is in RGB format
17
 
18
- # # Convert PIL image to bytes
19
- # # with io.BytesIO() as output:
20
- # # pil_image.save(output, format='JPEG')
21
- # # image_bytes = output.getvalue()
22
 
23
- # results = model.predict(image_array)
24
- # for result in results:
25
- # boxes = result.boxes.xyxy
26
- # conf = result.boxes.conf
27
- # cls = result.boxes.cls
28
- # obj_info = []
29
- # for i, bbox in enumerate(boxes):
30
- # label = result.names[int(cls[i])]
31
- # obj_info.append({
32
- # "Object": i+1,
33
- # "Label": label,
34
- # "Confidence": conf[i],
35
- # "Bounding Box": bbox
36
- # })
37
- # render = render_result(model=model, image=image, result=results[0])
38
- # if label:
39
- # print(label)
40
- # render.show()
41
- # return label
42
 
43
 
44
- from PIL import Image
45
- from io import BytesIO
46
 
47
- # Load model directly
48
- from transformers import AutoImageProcessor, AutoModelForObjectDetection
49
 
50
- processor = AutoImageProcessor.from_pretrained("savioratharv/pothole_detection")
51
- model = AutoModelForObjectDetection.from_pretrained("savioratharv/pothole_detection")
52
 
53
- # Function to predict if an image contains a pothole
54
- def predict_pothole(image_url):
55
- image = Image.open(BytesIO(image_url))
56
- inputs = processor(images=image, return_tensors="pt")
57
 
58
- # Perform inference
59
- outputs = model(**inputs)
60
- logits = outputs.logits
61
- probabilities = logits.softmax(dim=1)
62
 
63
- # Get predicted class (0: No pothole, 1: Pothole)
64
- predicted_class = probabilities.argmax().item()
65
- confidence = probabilities[0, predicted_class].item()
66
 
67
- return predicted_class
68
 
 
1
+ from ultralyticsplus import YOLO, render_result
2
+ from PIL import Image
3
+ from io import BytesIO
4
+ import numpy as np
5
 
6
+ def load_model(image):
7
+ # image_bytes = image.content
8
+ model = YOLO('keremberke/yolov8n-pothole-segmentation')
9
+ model.overrides['conf'] = 0.25
10
+ model.overrides['iou'] = 0.45
11
+ model.overrides['agnostic_nms'] = False
12
+ model.overrides['max_det'] = 1000
13
 
14
+ # Load image using PIL
15
+ image = Image.open(BytesIO(image))
16
+ image_array = np.array(image)
17
+ # pil_image = pil_image.convert("RGB") # Ensure image is in RGB format
18
 
19
+ # Convert PIL image to bytes
20
+ # with io.BytesIO() as output:
21
+ # pil_image.save(output, format='JPEG')
22
+ # image_bytes = output.getvalue()
23
 
24
+ results = model.predict(image_array)
25
+ for result in results:
26
+ boxes = result.boxes.xyxy
27
+ conf = result.boxes.conf
28
+ cls = result.boxes.cls
29
+ obj_info = []
30
+ for i, bbox in enumerate(boxes):
31
+ label = result.names[int(cls[i])]
32
+ obj_info.append({
33
+ "Object": i+1,
34
+ "Label": label,
35
+ "Confidence": conf[i],
36
+ "Bounding Box": bbox
37
+ })
38
+ render = render_result(model=model, image=image, result=results[0])
39
+ if label:
40
+ print(label)
41
+ render.show()
42
+ return label
43
 
44
 
45
+ # from PIL import Image
46
+ # from io import BytesIO
47
 
48
+ # # Load model directly
49
+ # from transformers import AutoImageProcessor, AutoModelForObjectDetection
50
 
51
+ # processor = AutoImageProcessor.from_pretrained("savioratharv/pothole_detection")
52
+ # model = AutoModelForObjectDetection.from_pretrained("savioratharv/pothole_detection")
53
 
54
+ # # Function to predict if an image contains a pothole
55
+ # def predict_pothole(image_url):
56
+ # image = Image.open(BytesIO(image_url))
57
+ # inputs = processor(images=image, return_tensors="pt")
58
 
59
+ # # Perform inference
60
+ # outputs = model(**inputs)
61
+ # logits = outputs.logits
62
+ # probabilities = logits.softmax(dim=1)
63
 
64
+ # # Get predicted class (0: No pothole, 1: Pothole)
65
+ # predicted_class = probabilities.argmax().item()
66
+ # confidence = probabilities[0, predicted_class].item()
67
 
68
+ # return predicted_class
69
 
requirements.txt CHANGED
@@ -8,5 +8,4 @@ python-multipart
8
  certifi
9
  firebase_admin
10
  pillow
11
- transformers
12
- pytorch-lightning
 
8
  certifi
9
  firebase_admin
10
  pillow
11
+ ultralyticsplus