Upload app.py

app.py

@@ -478,25 +478,33 @@ def generate_dataset_ui(num_samples, image_size):
         import os
         import json
         import random
-        from PIL import Image
-        import matplotlib.pyplot as plt
-        from matplotlib.patches import Rectangle
-        import io
+        from PIL import Image, ImageDraw
         
         output_dir = "./dataset"
         os.makedirs(output_dir, exist_ok=True)
         os.makedirs(os.path.join(output_dir, "images"), exist_ok=True)
         
-        bg_color = "#1a1a2e"
-        bullish_color = "#00ff88"
-        bearish_color = "#ff4466"
         num_candles = 20
+        labels = {}
+        
+        # Colors
+        bg_color = (26, 26, 46)
+        bullish_color = (0, 255, 136)
+        bearish_color = (255, 68, 102)
+        
+        patterns = ["bullish", "bearish", "sideways"]
+        volatilities = {"low": 1.0, "medium": 3.0, "high": 6.0}
         
-        def generate_candles(pattern, vol):
+        for i in range(int(num_samples)):
+            # Generate candle data
+            pattern = random.choice(patterns)
+            vol_name = random.choice(list(volatilities.keys()))
+            vol = volatilities[vol_name]
+            
             candles = []
             price = 100 if pattern != "bearish" else 150
             
-            for i in range(num_candles):
+            for j in range(num_candles):
                 if pattern == "bullish":
                     trend = random.uniform(0.5, 2.0)
                     o = price + random.gauss(0, vol)
@@ -505,7 +513,7 @@ def generate_dataset_ui(num_samples, image_size):
                     trend = random.uniform(0.5, 2.0)
                     o = price + random.gauss(0, vol)
                     c = o - random.uniform(0, vol*2) - trend
-                else:  # sideways
+                else:
                     o = price + random.gauss(0, vol)
                     c = o + random.gauss(0, vol)
                 
@@ -513,63 +521,62 @@ def generate_dataset_ui(num_samples, image_size):
                 l = min(o, c) - random.uniform(0, vol)
                 candles.append({"o": o, "h": h, "l": l, "c": c})
                 price = c
-            return candles
-        
-        def render(candles):
-            fig, ax = plt.subplots(figsize=(image_size/100, image_size/100), dpi=100)
-            fig.patch.set_facecolor(bg_color)
-            ax.set_facecolor(bg_color)
+            
+            # Render with PIL (no matplotlib)
+            img = Image.new('RGB', (image_size, image_size), bg_color)
+            draw = ImageDraw.Draw(img)
             
             highs = [c["h"] for c in candles]
             lows = [c["l"] for c in candles]
-            price_min, price_max = min(lows)*0.98, max(highs)*1.02
-            
-            for i, c in enumerate(candles):
-                color = bullish_color if c["c"] >= c["o"] else bearish_color
-                ax.plot([i, i], [c["l"], c["h"]], color=color, linewidth=1)
-                body_bottom = min(c["o"], c["c"])
-                body_height = abs(c["c"] - c["o"]) or 0.1
-                rect = Rectangle((i-0.3, body_bottom), 0.6, body_height, facecolor=color)
-                ax.add_patch(rect)
-            
-            ax.set_xlim(-1, len(candles))
-            ax.set_ylim(price_min, price_max)
-            ax.axis("off")
+            price_min, price_max = min(lows), max(highs)
+            price_range = price_max - price_min or 1
             
-            buf = io.BytesIO()
-            plt.savefig(buf, format="png", facecolor=bg_color, bbox_inches="tight", pad_inches=0.1)
-            plt.close(fig)
-            buf.seek(0)
-            
-            img = Image.open(buf).convert("RGB")
-            return img.resize((image_size, image_size), Image.Resampling.LANCZOS)
-        
-        patterns = ["bullish", "bearish", "sideways"]
-        volatilities = {"low": 1.0, "medium": 3.0, "high": 6.0}
-        labels = {}
-        
-        for i in range(int(num_samples)):
-            pattern = random.choice(patterns)
-            vol_name = random.choice(list(volatilities.keys()))
-            vol = volatilities[vol_name]
+            padding = 10
+            chart_width = image_size - 2 * padding
+            chart_height = image_size - 2 * padding
+            candle_width = chart_width // (num_candles + 2)
             
-            candles = generate_candles(pattern, vol)
-            img = render(candles)
+            for j, c in enumerate(candles):
+                x = padding + (j + 1) * candle_width
+                
+                # Scale prices to pixels
+                def scale_y(p):
+                    return padding + chart_height - int((p - price_min) / price_range * chart_height)
+                
+                y_open = scale_y(c["o"])
+                y_close = scale_y(c["c"])
+                y_high = scale_y(c["h"])
+                y_low = scale_y(c["l"])
+                
+                color = bullish_color if c["c"] >= c["o"] else bearish_color
+                
+                # Draw wick
+                draw.line([(x, y_high), (x, y_low)], fill=color, width=1)
+                
+                # Draw body
+                body_top = min(y_open, y_close)
+                body_bottom = max(y_open, y_close)
+                if body_bottom - body_top < 2:
+                    body_bottom = body_top + 2
+                draw.rectangle(
+                    [(x - candle_width//3, body_top), (x + candle_width//3, body_bottom)],
+                    fill=color
+                )
             
+            # Save image
             filename = f"chart_{i:06d}.png"
             img.save(os.path.join(output_dir, "images", filename))
             labels[filename] = f"{pattern} trend {vol_name} volatility"
-            
-            if i % 500 == 0:
-                print(f"Generated {i}/{num_samples}")
         
+        # Save labels
         with open(os.path.join(output_dir, "labels.json"), "w") as f:
             json.dump(labels, f)
         
-        return f"✅ Generated {num_samples} samples in ./dataset"
+        return f"✅ Generated {int(num_samples)} samples in ./dataset"
     
     except Exception as e:
-        return f"❌ Failed: {str(e)}"
+        import traceback
+        return f"❌ Failed: {str(e)}\n{traceback.format_exc()}"
 
 
 # ============== Gradio Interface ==============