Upload app.py

app.py

@@ -436,6 +436,106 @@ def load_model(checkpoint_path=None):
     return True
 
 
+def generate_dataset_ui(num_samples, image_size):
+    """Generate training dataset."""
+    try:
+        import os
+        import json
+        import random
+        from PIL import Image
+        import matplotlib.pyplot as plt
+        from matplotlib.patches import Rectangle
+        import io
+        
+        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
+        
+        def generate_candles(pattern, vol):
+            candles = []
+            price = 100 if pattern != "bearish" else 150
+            
+            for i in range(num_candles):
+                if pattern == "bullish":
+                    trend = random.uniform(0.5, 2.0)
+                    o = price + random.gauss(0, vol)
+                    c = o + random.uniform(0, vol*2) + trend
+                elif pattern == "bearish":
+                    trend = random.uniform(0.5, 2.0)
+                    o = price + random.gauss(0, vol)
+                    c = o - random.uniform(0, vol*2) - trend
+                else:  # sideways
+                    o = price + random.gauss(0, vol)
+                    c = o + random.gauss(0, vol)
+                
+                h = max(o, c) + random.uniform(0, vol)
+                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)
+            
+            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")
+            
+            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]
+            
+            candles = generate_candles(pattern, vol)
+            img = render(candles)
+            
+            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}")
+        
+        with open(os.path.join(output_dir, "labels.json"), "w") as f:
+            json.dump(labels, f)
+        
+        return f"✅ Generated {num_samples} samples in ./dataset"
+    
+    except Exception as e:
+        return f"❌ Failed: {str(e)}"
+
+
 # ============== Gradio Interface ==============
 
 def generate_chart(prompt, num_steps, guidance_scale, seed):
@@ -592,6 +692,30 @@ def create_demo():
         """)
         
         with gr.Tabs():
+            # Data Generation Tab
+            with gr.TabItem("📊 Generate Data"):
+                gr.Markdown("""
+                ### Generate Training Dataset
+                
+                Create synthetic candlestick chart images for training.
+                **Run this first before training!**
+                """)
+                
+                with gr.Row():
+                    with gr.Column():
+                        num_samples = gr.Slider(1000, 50000, value=10000, step=1000, label="Number of Samples")
+                        data_image_size = gr.Slider(64, 256, value=128, step=32, label="Image Size")
+                        generate_data_btn = gr.Button("📊 Generate Dataset", variant="primary")
+                    
+                    with gr.Column():
+                        data_status = gr.Textbox(label="Status", lines=5, interactive=False)
+                
+                generate_data_btn.click(
+                    generate_dataset_ui,
+                    inputs=[num_samples, data_image_size],
+                    outputs=[data_status]
+                )
+            
             # Generation Tab
             with gr.TabItem("🎨 Generate"):
                 with gr.Row():