Update colab.py

colab.py

@@ -194,62 +194,95 @@ class DanceGenerator:
     def _create_enhanced_dance_frame(self, pose_array, frame_size, add_effects=True):
         """Create enhanced visualization frame with effects"""
         height, width = frame_size
-        # Change background from black to light gray for better visibility
-        frame = np.ones((height, width, 3), dtype=np.uint8) * 240  # Light gray background
+        # Create transparent background
+        frame = np.zeros((height, width, 3), dtype=np.uint8)  # Black background
         
         # Convert coordinates
         points = (pose_array[:, :2] * [width, height]).astype(int)
         
-        # Draw enhanced skeleton with thicker lines and more visible colors
+        # Draw enhanced skeleton with neon effects
         connections = self._get_pose_connections()
-        for connection in connections:
-            start_idx, end_idx = connection
-            if start_idx < len(points) and end_idx < len(points):
-                if add_effects:
-                    self._draw_glowing_line(
-                        frame,
-                        points[start_idx],
-                        points[end_idx],
-                        (0, 100, 255),  # Orange color for skeleton
-                        thickness=4
-                    )
-                else:
-                    cv2.line(frame,
-                            tuple(points[start_idx]),
-                            tuple(points[end_idx]),
-                            (0, 100, 255), 4)
-
-        # Draw enhanced joints with larger radius
-        for point in points:
+        
+        # Define body parts and their colors
+        body_parts = {
+            'spine': [(11, 23), (23, 24), (11, 12)],  # Torso
+            'right_arm': [(11, 13), (13, 15)],  # Right arm
+            'left_arm': [(12, 14), (14, 16)],   # Left arm
+            'right_leg': [(23, 25), (25, 27), (27, 29), (29, 31)],  # Right leg
+            'left_leg': [(24, 26), (26, 28), (28, 30), (30, 32)],   # Left leg
+            'face': [(0, 1), (1, 2), (2, 3), (3, 7), (0, 4), (4, 5), (5, 6), (6, 8)]  # Face
+        }
+        
+        colors = {
+            'spine': (0, 255, 255),    # Cyan
+            'right_arm': (0, 255, 0),  # Green
+            'left_arm': (0, 255, 0),   # Green
+            'right_leg': (255, 0, 255), # Magenta
+            'left_leg': (255, 0, 255),  # Magenta
+            'face': (255, 255, 0)      # Yellow
+        }
+        
+        # Draw body parts with glow effects
+        for part, connections_list in body_parts.items():
+            color = colors[part]
+            for connection in connections_list:
+                start_idx, end_idx = connection
+                if start_idx < len(points) and end_idx < len(points):
+                    if add_effects:
+                        self._draw_glowing_line(
+                            frame,
+                            points[start_idx],
+                            points[end_idx],
+                            color,
+                            thickness=3
+                        )
+                    else:
+                        cv2.line(frame,
+                                tuple(points[start_idx]),
+                                tuple(points[end_idx]),
+                                color, 3)
+
+        # Draw enhanced joints with glow
+        for i, point in enumerate(points):
             if add_effects:
-                self._draw_glowing_point(frame, point, (255, 0, 0), radius=6)  # Blue joints
+                # Different colors for different body parts
+                if i in [0,1,2,3,4,5,6,7,8]:  # Face points
+                    color = (255, 255, 0)  # Yellow
+                elif i in [11,12,23,24]:  # Torso points
+                    color = (0, 255, 255)  # Cyan
+                elif i in [13,14,15,16]:  # Arms points
+                    color = (0, 255, 0)    # Green
+                else:  # Legs points
+                    color = (255, 0, 255)  # Magenta
+                    
+                self._draw_glowing_point(frame, point, color, radius=4)
             else:
-                cv2.circle(frame, tuple(point), 6, (255, 0, 0), -1)
+                cv2.circle(frame, tuple(point), 4, (255, 255, 255), -1)
 
         return frame
 
-    def _draw_glowing_line(self, frame, start, end, color, thickness=4):
-        """Draw a line with enhanced glow effect"""
+    def _draw_glowing_line(self, frame, start, end, color, thickness=3):
+        """Draw a line with enhanced neon glow effect"""
         # Draw outer glow
         for i in range(3):
-            alpha = 0.5 - i * 0.15
-            thick = thickness + (i * 4)
+            alpha = 0.3 - i * 0.1
+            thick = thickness + (i * 2)
+            blur_color = tuple([int(c * alpha) for c in color])
             cv2.line(frame, tuple(start), tuple(end),
-                    tuple([int(c * alpha) for c in color]),
-                    thick)
+                    blur_color, thick)
         
         # Draw main line
         cv2.line(frame, tuple(start), tuple(end), color, thickness)
 
-    def _draw_glowing_point(self, frame, point, color, radius=6):
-        """Draw a point with enhanced glow effect"""
+    def _draw_glowing_point(self, frame, point, color, radius=4):
+        """Draw a point with enhanced neon glow effect"""
         # Draw outer glow
         for i in range(3):
-            alpha = 0.5 - i * 0.15
-            r = radius + (i * 3)
+            alpha = 0.3 - i * 0.1
+            r = radius + (i * 2)
+            blur_color = tuple([int(c * alpha) for c in color])
             cv2.circle(frame, tuple(point), r,
-                      tuple([int(c * alpha) for c in color]),
-                      -1)
+                      blur_color, -1)
         
         # Draw main point
         cv2.circle(frame, tuple(point), radius, color, -1)