capradeepgujaran commited on
Commit
18cd948
1 Parent(s): b122109

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +69 -125
app.py CHANGED
@@ -7,6 +7,7 @@ from PIL import Image as PILImage
7
  import io
8
  import os
9
  import base64
 
10
 
11
  def create_monitor_interface():
12
  api_key = os.getenv("GROQ_API_KEY")
@@ -15,14 +16,27 @@ def create_monitor_interface():
15
  def __init__(self):
16
  self.client = Groq()
17
  self.model_name = "llama-3.2-90b-vision-preview"
18
- self.max_image_size = (800, 800)
19
- self.colors = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (255, 255, 0), (255, 0, 255)]
20
 
 
 
 
 
 
 
 
 
 
 
 
 
 
21
  def analyze_frame(self, frame: np.ndarray) -> str:
22
  if frame is None:
23
- return ""
24
-
25
- # Convert image
26
  if len(frame.shape) == 2:
27
  frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
28
  elif len(frame.shape) == 3 and frame.shape[2] == 4:
@@ -31,11 +45,12 @@ def create_monitor_interface():
31
  frame = self.resize_image(frame)
32
  frame_pil = PILImage.fromarray(frame)
33
 
 
34
  buffered = io.BytesIO()
35
  frame_pil.save(buffered,
36
- format="JPEG",
37
- quality=95, # Increased quality
38
- optimize=True)
39
  img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
40
  image_url = f"data:image/jpeg;base64,{img_base64}"
41
 
@@ -43,36 +58,14 @@ def create_monitor_interface():
43
  completion = self.client.chat.completions.create(
44
  model=self.model_name,
45
  messages=[
46
- {
47
- "role": "system",
48
- "content": """You are a construction site safety expert specializing in ergonomics and workplace safety.
49
- Analyze images for:
50
- 1. Worker posture and ergonomic risks
51
- 2. PPE usage and compliance
52
- 3. Tool and equipment safety
53
- 4. Environmental hazards
54
- 5. Working position and technique"""
55
- },
56
  {
57
  "role": "user",
58
  "content": [
59
  {
60
  "type": "text",
61
- "text": """Carefully analyze this construction worker's position and environment. Look for:
62
-
63
- 1. Ergonomic issues (kneeling position, back posture, repetitive motions)
64
- 2. PPE compliance (knee pads, gloves, appropriate footwear)
65
- 3. Working technique and body mechanics
66
- 4. Surrounding hazards or risks
67
-
68
- For each issue identified, format your response as:
69
- - <location>position:specific safety concern and recommendation</location>
70
-
71
- For example:
72
- - <location>center:Worker kneeling without knee protection, risking joint injury. Recommend knee pads.</location>
73
- - <location>bottom:Improper back posture while working, potential for strain. Should maintain straight back.</location>
74
-
75
- Be specific about each safety concern you observe."""
76
  },
77
  {
78
  "type": "image_url",
@@ -81,127 +74,78 @@ def create_monitor_interface():
81
  }
82
  }
83
  ]
 
 
 
 
84
  }
85
  ],
86
- temperature=0.3, # Lowered temperature for more focused analysis
87
- max_tokens=500,
88
- stream=False
 
 
89
  )
90
-
91
- response = completion.choices[0].message.content
92
- print(f"Raw response: {response}") # For debugging
93
- return response
94
-
95
  except Exception as e:
96
- print(f"Analysis error: {str(e)}")
97
- return ""
98
-
99
- def resize_image(self, image):
100
- height, width = image.shape[:2]
101
- if height > self.max_image_size[1] or width > self.max_image_size[0]:
102
- aspect = width / height
103
- if width > height:
104
- new_width = self.max_image_size[0]
105
- new_height = int(new_width / aspect)
106
- else:
107
- new_height = self.max_image_size[1]
108
- new_width = int(new_height * aspect)
109
- return cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
110
- return image
111
-
112
- def get_region_coordinates(self, position: str, image_shape: tuple) -> tuple:
113
- height, width = image_shape[:2]
114
- regions = {
115
- 'top-left': (0, 0, width//3, height//3),
116
- 'top': (width//3, 0, 2*width//3, height//3),
117
- 'top-right': (2*width//3, 0, width, height//3),
118
- 'left': (0, height//3, width//3, 2*height//3),
119
- 'center': (width//3, height//3, 2*width//3, 2*height//3),
120
- 'right': (2*width//3, height//3, width, 2*height//3),
121
- 'bottom-left': (0, 2*height//3, width//3, height),
122
- 'bottom': (width//3, 2*height//3, 2*width//3, height),
123
- 'bottom-right': (2*width//3, 2*height//3, width, height)
124
- }
125
-
126
- # Try to match the position with regions
127
- matched_region = None
128
- max_match_length = 0
129
- position_lower = position.lower()
130
-
131
- for region_name in regions:
132
- if region_name in position_lower:
133
- if len(region_name) > max_match_length:
134
- matched_region = region_name
135
- max_match_length = len(region_name)
136
-
137
- if matched_region:
138
- return regions[matched_region]
139
- return regions['center']
140
 
141
  def draw_observations(self, image, observations):
142
  height, width = image.shape[:2]
143
  font = cv2.FONT_HERSHEY_SIMPLEX
144
- font_scale = 0.6
145
  thickness = 2
146
 
 
147
  for idx, obs in enumerate(observations):
148
  color = self.colors[idx % len(self.colors)]
149
 
150
- parts = obs.split(':')
151
- if len(parts) >= 2:
152
- position = parts[0]
153
- description = ':'.join(parts[1:])
154
-
155
- x1, y1, x2, y2 = self.get_region_coordinates(position, image.shape)
156
-
157
- # Draw rectangle
158
- cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
159
-
160
- # Add label with background
161
- label = description[:50] + "..." if len(description) > 50 else description
162
- label_size = cv2.getTextSize(label, font, font_scale, thickness)[0]
163
-
164
- label_x = max(0, min(x1, width - label_size[0]))
165
- label_y = max(20, y1 - 5)
166
-
167
- cv2.rectangle(image, (label_x, label_y - 20),
168
- (label_x + label_size[0], label_y), color, -1)
169
- cv2.putText(image, label, (label_x, label_y - 5),
170
- font, font_scale, (255, 255, 255), thickness)
171
 
172
  return image
173
 
174
  def process_frame(self, frame: np.ndarray) -> tuple[np.ndarray, str]:
175
  if frame is None:
176
  return None, "No image provided"
177
-
178
  analysis = self.analyze_frame(frame)
179
- print(f"Analysis received: {analysis}") # Debug print
180
 
 
181
  observations = []
182
  for line in analysis.split('\n'):
183
  line = line.strip()
184
  if line.startswith('-'):
 
185
  if '<location>' in line and '</location>' in line:
186
  start = line.find('<location>') + len('<location>')
187
  end = line.find('</location>')
188
- observation = line[start:end].strip()
189
- if observation and ':' in observation:
190
- observations.append(observation)
191
-
192
- print(f"Parsed observations: {observations}") # Debug print
193
-
194
- display_frame = frame.copy()
195
- if observations:
196
- annotated_frame = self.draw_observations(display_frame, observations)
197
- return annotated_frame, analysis
198
 
199
- # If no observations were found but we got some analysis
200
- if analysis and not analysis.isspace():
201
- return display_frame, analysis
202
 
203
- return display_frame, "Please try again - no safety analysis was generated."
204
 
 
205
  monitor = SafetyMonitor()
206
 
207
  with gr.Blocks() as demo:
@@ -209,9 +153,9 @@ def create_monitor_interface():
209
 
210
  with gr.Row():
211
  input_image = gr.Image(label="Upload Image")
212
- output_image = gr.Image(label="Analysis Results")
213
 
214
- analysis_text = gr.Textbox(label="Safety Analysis", lines=5)
215
 
216
  def analyze_image(image):
217
  if image is None:
 
7
  import io
8
  import os
9
  import base64
10
+ import random
11
 
12
  def create_monitor_interface():
13
  api_key = os.getenv("GROQ_API_KEY")
 
16
  def __init__(self):
17
  self.client = Groq()
18
  self.model_name = "llama-3.2-90b-vision-preview"
19
+ self.max_image_size = (640, 640) # Increased size for better visibility
20
+ self.colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]
21
 
22
+ def resize_image(self, image):
23
+ height, width = image.shape[:2]
24
+ aspect = width / height
25
+
26
+ if width > height:
27
+ new_width = min(self.max_image_size[0], width)
28
+ new_height = int(new_width / aspect)
29
+ else:
30
+ new_height = min(self.max_image_size[1], height)
31
+ new_width = int(new_height * aspect)
32
+
33
+ return cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
34
+
35
  def analyze_frame(self, frame: np.ndarray) -> str:
36
  if frame is None:
37
+ return "No frame received"
38
+
39
+ # Convert and resize image
40
  if len(frame.shape) == 2:
41
  frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
42
  elif len(frame.shape) == 3 and frame.shape[2] == 4:
 
45
  frame = self.resize_image(frame)
46
  frame_pil = PILImage.fromarray(frame)
47
 
48
+ # Convert to base64 with minimal quality
49
  buffered = io.BytesIO()
50
  frame_pil.save(buffered,
51
+ format="JPEG",
52
+ quality=30,
53
+ optimize=True)
54
  img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
55
  image_url = f"data:image/jpeg;base64,{img_base64}"
56
 
 
58
  completion = self.client.chat.completions.create(
59
  model=self.model_name,
60
  messages=[
 
 
 
 
 
 
 
 
 
 
61
  {
62
  "role": "user",
63
  "content": [
64
  {
65
  "type": "text",
66
+ "text": """Analyze this workplace image and describe each safety concern in this format:
67
+ - <location>Description</location>
68
+ Use one line per issue, starting with a dash and location in tags."""
 
 
 
 
 
 
 
 
 
 
 
 
69
  },
70
  {
71
  "type": "image_url",
 
74
  }
75
  }
76
  ]
77
+ },
78
+ {
79
+ "role": "assistant",
80
+ "content": ""
81
  }
82
  ],
83
+ temperature=0.1,
84
+ max_tokens=150,
85
+ top_p=1,
86
+ stream=False,
87
+ stop=None
88
  )
89
+ return completion.choices[0].message.content
 
 
 
 
90
  except Exception as e:
91
+ print(f"Detailed error: {str(e)}")
92
+ return f"Analysis Error: {str(e)}"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
93
 
94
  def draw_observations(self, image, observations):
95
  height, width = image.shape[:2]
96
  font = cv2.FONT_HERSHEY_SIMPLEX
97
+ font_scale = 0.5
98
  thickness = 2
99
 
100
+ # Generate random positions for each observation
101
  for idx, obs in enumerate(observations):
102
  color = self.colors[idx % len(self.colors)]
103
 
104
+ # Generate random box position
105
+ box_width = width // 3
106
+ box_height = height // 3
107
+ x = random.randint(0, width - box_width)
108
+ y = random.randint(0, height - box_height)
109
+
110
+ # Draw rectangle
111
+ cv2.rectangle(image, (x, y), (x + box_width, y + box_height), color, 2)
112
+
113
+ # Add label with background
114
+ label = obs[:40] + "..." if len(obs) > 40 else obs
115
+ label_size = cv2.getTextSize(label, font, font_scale, thickness)[0]
116
+ cv2.rectangle(image, (x, y - 20), (x + label_size[0], y), color, -1)
117
+ cv2.putText(image, label, (x, y - 5), font, font_scale, (255, 255, 255), thickness)
 
 
 
 
 
 
 
118
 
119
  return image
120
 
121
  def process_frame(self, frame: np.ndarray) -> tuple[np.ndarray, str]:
122
  if frame is None:
123
  return None, "No image provided"
124
+
125
  analysis = self.analyze_frame(frame)
126
+ display_frame = self.resize_image(frame.copy())
127
 
128
+ # Parse observations from the analysis
129
  observations = []
130
  for line in analysis.split('\n'):
131
  line = line.strip()
132
  if line.startswith('-'):
133
+ # Extract text between <location> tags if present
134
  if '<location>' in line and '</location>' in line:
135
  start = line.find('<location>') + len('<location>')
136
  end = line.find('</location>')
137
+ observation = line[end + len('</location>'):].strip()
138
+ else:
139
+ observation = line[1:].strip() # Remove the dash
140
+ if observation:
141
+ observations.append(observation)
 
 
 
 
 
142
 
143
+ # Draw observations on the image
144
+ annotated_frame = self.draw_observations(display_frame, observations)
 
145
 
146
+ return annotated_frame, analysis
147
 
148
+ # Create the main interface
149
  monitor = SafetyMonitor()
150
 
151
  with gr.Blocks() as demo:
 
153
 
154
  with gr.Row():
155
  input_image = gr.Image(label="Upload Image")
156
+ output_image = gr.Image(label="Annotated Results")
157
 
158
+ analysis_text = gr.Textbox(label="Detailed Analysis", lines=5)
159
 
160
  def analyze_image(image):
161
  if image is None: