Upload 3 files

best_model.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dd85c6f84bca45ea689fb0a0e402b5432a86bf2020046b52cf3b1e5aa19bf041
-size 17424939
+oid sha256:134657576082cf3606a40722ead3d5e5df9111178f74f17584263de4082c300c
+size 72353922

imageAI.py

@@ -1,22 +1,22 @@
 try:
-  import google.colab
-  IN_COLAB = True
-  from google.colab import drive,files
-  from google.colab import output
-  drive.mount('/gdrive')
-  Gbase="/gdrive/MyDrive/generate/"
-  cache_dir="/gdrive/MyDrive/hf/"
-  import sys
-  sys.path.append(Gbase)
+    import google.colab
+    IN_COLAB = True
+    from google.colab import drive, files
+    from google.colab import output
+    drive.mount('/gdrive')
+    Gbase = "/gdrive/MyDrive/generate/"
+    cache_dir = "/gdrive/MyDrive/hf/"
+    import sys
+    sys.path.append(Gbase)
 except:
-  IN_COLAB = False
-  Gbase="./"
-  cache_dir="./hf/"
+    IN_COLAB = False
+    Gbase = "./"
+    cache_dir = "./hf/"
 
-
-import cv2,os
+import cv2
+import os
 import numpy as np
-import random,string
+import random
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -28,36 +28,129 @@ print(f"Using device: {device}")
 
 IMAGE_SIZE = 64
 NUM_SAMPLES = 1000
-BATCH_SIZE = 4
+BATCH_SIZE = 200
 EPOCHS = 500
 LEARNING_RATE = 0.001
 
+def generate_sample(num_shapes=1):
+    image = np.zeros((IMAGE_SIZE, IMAGE_SIZE, 3), dtype=np.uint8)
+    instructions = []
+    
+    for _ in range(num_shapes):
+        shape = random.choice(['rectangle', 'circle', 'ellipse', 'polygon'])
+        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+        
+        if shape == 'rectangle':
+            start_point = (random.randint(0, IMAGE_SIZE - 10), random.randint(0, IMAGE_SIZE - 10))
+            end_point = (start_point[0] + random.randint(10, IMAGE_SIZE - start_point[0]),
+                         start_point[1] + random.randint(10, IMAGE_SIZE - start_point[1]))
+            cv2.rectangle(image, start_point, end_point, color, -1)
+            instructions.append(f"cv2.rectangle(image, {start_point}, {end_point}, {color}, -1)")
+        
+        elif shape == 'circle':
+            center = (random.randint(10, IMAGE_SIZE - 10), random.randint(10, IMAGE_SIZE - 10))
+            radius = random.randint(5, min(center[0], center[1], IMAGE_SIZE - center[0], IMAGE_SIZE - center[1]))
+            cv2.circle(image, center, radius, color, -1)
+            instructions.append(f"cv2.circle(image, {center}, {radius}, {color}, -1)")
+        
+        elif shape == 'ellipse':
+            center = (random.randint(10, IMAGE_SIZE - 10), random.randint(10, IMAGE_SIZE - 10))
+            axes = (random.randint(5, 30), random.randint(5, 30))
+            angle = random.randint(0, 360)
+            cv2.ellipse(image, center, axes, angle, 0, 360, color, -1)
+            instructions.append(f"cv2.ellipse(image, {center}, {axes}, {angle}, 0, 360, {color}, -1)")
+        
+        elif shape == 'polygon':
+            num_points = random.randint(3, 6)
+            points = np.array([(random.randint(0, IMAGE_SIZE), random.randint(0, IMAGE_SIZE)) for _ in range(num_points)], np.int32)
+            points = points.reshape((-1, 1, 2))
+            cv2.fillPoly(image, [points], color)
+            instructions.append(f"cv2.fillPoly(image, [{points.tolist()}], {color})")
+    
+    return {'image': image, 'instructions': instructions}
+
+def generate_dataset(NUM_SAMPLES=NUM_SAMPLES, maxNumShape=3):
+    dataset = []
+    for _ in range(NUM_SAMPLES):
+        num_shapes = random.randint(1, maxNumShape)
+        sample = generate_sample(num_shapes=num_shapes)
+        dataset.append(sample)
+    return dataset
+
+class ImageDataset(Dataset):
+    def __init__(self, dataset):
+        self.dataset = dataset
+    
+    def __len__(self):
+        return len(self.dataset)
+    
+    def __getitem__(self, idx):
+        sample = self.dataset[idx]
+        image = torch.FloatTensor(sample['image']).permute(2, 0, 1) / 255.0
+        return image, len(sample['instructions'])
+
+class SelfAttention(nn.Module):
+    def __init__(self, in_channels):
+        super(SelfAttention, self).__init__()
+        self.query = nn.Conv2d(in_channels, in_channels // 8, 1)
+        self.key = nn.Conv2d(in_channels, in_channels // 8, 1)
+        self.value = nn.Conv2d(in_channels, in_channels, 1)
+        self.gamma = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x):
+        batch_size, C, width, height = x.size()
+        
+        proj_query = self.query(x).view(batch_size, -1, width * height).permute(0, 2, 1)
+        proj_key = self.key(x).view(batch_size, -1, width * height)
+        energy = torch.bmm(proj_query, proj_key)
+        attention = F.softmax(energy, dim=-1)
+        
+        proj_value = self.value(x).view(batch_size, -1, width * height)
+        out = torch.bmm(proj_value, attention.permute(0, 2, 1))
+        out = out.view(batch_size, C, width, height)
+        
+        out = self.gamma * out + x
+        return out
+
 class SimpleModel(nn.Module):
     def __init__(self, path=None):
         super(SimpleModel, self).__init__()
-        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
-        self.bn1 = nn.BatchNorm2d(32)
-        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
-        self.bn2 = nn.BatchNorm2d(64)
-        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
-        self.bn3 = nn.BatchNorm2d(128)
-        self.pool = nn.MaxPool2d(2, 2)
-        self.fc1 = nn.Linear(128 * 8 * 8, 512)
-        self.fc2 = nn.Linear(512, 128)
-        self.fc3 = nn.Linear(128, 1)
+        
+        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
+        self.bn2 = nn.BatchNorm2d(128)
+        self.conv3 = nn.Conv2d(128, 256, 3, padding=1)
+        self.bn3 = nn.BatchNorm2d(256)
+        self.conv4 = nn.Conv2d(256, 256, 3, padding=1)
+        self.bn4 = nn.BatchNorm2d(256)
+        
+        self.pool = nn.AdaptiveAvgPool2d((8, 8))
+        
+        self.attention = SelfAttention(256)
+        
+        self.fc1 = nn.Linear(256 * 8 * 8, 1024)
+        self.fc2 = nn.Linear(1024, 256)
+        self.fc3 = nn.Linear(256, 1)
+        
         self.dropout = nn.Dropout(0.5)
         
         if path and os.path.exists(path):
             self.load_state_dict(torch.load(path, map_location=device))
     
     def forward(self, x):
-        x = self.pool(F.leaky_relu(self.bn1(self.conv1(x))))
-        x = self.pool(F.leaky_relu(self.bn2(self.conv2(x))))
-        x = self.pool(F.leaky_relu(self.bn3(self.conv3(x))))
-        x = x.view(-1, 128 * 8 * 8)
-        x = F.leaky_relu(self.fc1(x))
+        x = self.pool(F.mish(self.bn1(self.conv1(x))))
+        x = self.pool(F.mish(self.bn2(self.conv2(x))))
+        x = F.mish(self.bn3(self.conv3(x)))
+        x = F.mish(self.bn4(self.conv4(x)))
+        
+        x = self.attention(x)
+        
+        x = self.pool(x)
+        x = x.view(-1, 256 * 8 * 8)
+        x = F.mish(self.fc1(x))
         x = self.dropout(x)
-        x = F.leaky_relu(self.fc2(x))
+        x = F.mish(self.fc2(x))
         x = self.dropout(x)
         x = self.fc3(x)
         return x
@@ -66,43 +159,147 @@ class SimpleModel(nn.Module):
         self.eval()
         with torch.no_grad():
             if isinstance(image, str) and os.path.isfile(image):
-                # 如果輸入是圖片文件路徑
-                img = cv2.imread(image, cv2.IMREAD_GRAYSCALE)
+                img = cv2.imread(image)
                 img = cv2.resize(img, (IMAGE_SIZE, IMAGE_SIZE))
             elif isinstance(image, np.ndarray):
-                # 如果輸入是 numpy 數組
-                if image.ndim == 3:
-                    img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-                else:
-                    img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+                if image.ndim == 2:
+                    img = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
                 img = cv2.resize(img, (IMAGE_SIZE, IMAGE_SIZE))
             else:
                 raise ValueError("Input should be an image file path or a numpy array")
             
-            img_tensor = torch.FloatTensor(img).unsqueeze(0).unsqueeze(0) / 255.0
+            img_tensor = torch.FloatTensor(img).permute(2, 0, 1).unsqueeze(0) / 255.0
             img_tensor = img_tensor.to(device)
             output = self(img_tensor).item()
             
-            # 將輸出四捨五入到最接近的整數
             num_instructions = round(output)
             
-            # 生成相應數量的繪圖指令
             instructions = []
             for _ in range(num_instructions):
-                shape = random.choice(['line', 'rectangle', 'circle', 'ellipse', 'polygon'])
-                if shape == 'line':
-                    instructions.append(f"cv2.line(image, {(random.randint(0, IMAGE_SIZE), random.randint(0, IMAGE_SIZE))}, {(random.randint(0, IMAGE_SIZE), random.randint(0, IMAGE_SIZE))}, {random.randint(0, 255)}, {random.randint(1, 3)})")
-                elif shape == 'rectangle':
-                    instructions.append(f"cv2.rectangle(image, {(random.randint(0, IMAGE_SIZE-10), random.randint(0, IMAGE_SIZE-10))}, {(random.randint(10, IMAGE_SIZE), random.randint(10, IMAGE_SIZE))}, {random.randint(0, 255)}, {random.randint(1, 3)})")
+                shape = random.choice(['rectangle', 'circle', 'ellipse', 'polygon'])
+                color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+                if shape == 'rectangle':
+                    instructions.append(f"cv2.rectangle(image, {(random.randint(0, IMAGE_SIZE-10), random.randint(0, IMAGE_SIZE-10))}, {(random.randint(10, IMAGE_SIZE), random.randint(10, IMAGE_SIZE))}, {color}, -1)")
                 elif shape == 'circle':
-                    instructions.append(f"cv2.circle(image, {(random.randint(10, IMAGE_SIZE-10), random.randint(10, IMAGE_SIZE-10))}, {random.randint(5, 30)}, {random.randint(0, 255)}, {random.randint(1, 3)})")
+                    instructions.append(f"cv2.circle(image, {(random.randint(10, IMAGE_SIZE-10), random.randint(10, IMAGE_SIZE-10))}, {random.randint(5, 30)}, {color}, -1)")
                 elif shape == 'ellipse':
-                    instructions.append(f"cv2.ellipse(image, {(random.randint(10, IMAGE_SIZE-10), random.randint(10, IMAGE_SIZE-10))}, {(random.randint(5, 30), random.randint(5, 30))}, {random.randint(0, 360)}, 0, 360, {random.randint(0, 255)}, {random.randint(1, 3)})")
+                    instructions.append(f"cv2.ellipse(image, {(random.randint(10, IMAGE_SIZE-10), random.randint(10, IMAGE_SIZE-10))}, {(random.randint(5, 30), random.randint(5, 30))}, {random.randint(0, 360)}, 0, 360, {color}, -1)")
                 elif shape == 'polygon':
                     num_points = random.randint(3, 6)
                     points = [(random.randint(0, IMAGE_SIZE), random.randint(0, IMAGE_SIZE)) for _ in range(num_points)]
-                    instructions.append(f"cv2.polylines(image, [np.array({points})], True, {random.randint(0, 255)}, {random.randint(1, 3)})")
-              
+                    instructions.append(f"cv2.fillPoly(image, [np.array({points})], {color})")
             
             return instructions
 
+def train(model, train_loader, optimizer, criterion):
+    model.train()
+    total_loss = 0
+    correct_predictions = 0
+    total_predictions = 0
+    for batch_idx, (data, target) in enumerate(train_loader):
+        data, target = data.to(device), target.float().to(device)
+        optimizer.zero_grad()
+        output = model(data).squeeze()
+        loss = criterion(output, target)
+        loss.backward()
+        optimizer.step()
+        total_loss += loss.item()
+        
+        # Count correct predictions
+        predicted = torch.round(output)
+        correct_predictions += (predicted == target).sum().item()
+        total_predictions += target.size(0)
+        
+        if batch_idx % 3000 == 0:
+            print(f'Train Batch {batch_idx}/{len(train_loader)} Loss: {loss.item():.6f}')
+    
+    accuracy = correct_predictions / total_predictions
+    return total_loss / len(train_loader), accuracy
+
+def test(model, test_loader, criterion, print_predictions=False):
+    model.eval()
+    test_loss = 0
+    correct_predictions = 0
+    total_predictions = 0
+    all_predictions = []
+    all_targets = []
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.float().to(device)
+            output = model(data).squeeze()
+            test_loss += criterion(output, target).item()
+            
+            predicted = torch.round(output)
+            correct_predictions += (predicted == target).sum().item()
+            total_predictions += target.size(0)
+            
+            all_predictions.extend(output.cpu().numpy())
+            all_targets.extend(target.cpu().numpy())
+    
+    test_loss /= len(test_loader)
+    accuracy = correct_predictions / total_predictions
+    print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {accuracy:.4f}')
+    
+    if print_predictions:
+        print("Sample predictions:")
+        for pred, targ in zip(all_predictions[:10], all_targets[:10]):
+            print(f"Prediction: {pred:.2f}, Target: {targ:.2f}")
+    
+    return test_loss, accuracy, all_predictions, all_targets
+
+def train1(NUM_SAMPLES=NUM_SAMPLES, maxNumShape=1, EPOCHS=EPOCHS):
+    model = SimpleModel(path=os.path.join(Gbase, 'best_model.pth')).to(device)
+    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
+    
+    optimizer_path = os.path.join(Gbase, 'optimizer.pth')
+    if os.path.exists(optimizer_path):
+        print("Loading optimizer state...")
+        optimizer.load_state_dict(torch.load(optimizer_path, map_location=device))
+    
+    criterion = nn.MSELoss()
+    
+    seed = 618 * 382 * 33
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+
+    dataset = generate_dataset(NUM_SAMPLES=NUM_SAMPLES, maxNumShape=maxNumShape)
+    train_size = int(0.8 * len(dataset))
+    train_dataset = ImageDataset(dataset[:train_size])
+    test_dataset = ImageDataset(dataset[train_size:])
+
+    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
+    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+    best_loss = float('inf')
+    best_accuracy = 0
+
+    for epoch in range(EPOCHS):
+        print(f'Epoch {epoch+1}/{EPOCHS}')
+        train_loss, train_accuracy = train(model, train_loader, optimizer, criterion)
+        test_loss, test_accuracy, predictions, targets = test(model, test_loader, criterion, print_predictions=True)
+        
+        print(f'Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.4f}')
+        print(f'Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.4f}')
+        
+        if test_accuracy > best_accuracy or (test_accuracy == best_accuracy and test_loss < best_loss):
+            best_accuracy = test_accuracy
+            best_loss = test_loss
+            torch.save(model.state_dict(), os.path.join(Gbase, 'best_model.pth') ,_use_new_zipfile_serialization=False )
+            torch.save(optimizer.state_dict(), os.path.join(Gbase, 'optimizer.pth') ,_use_new_zipfile_serialization=False)
+            print(f"New best model saved with test accuracy: {best_accuracy:.4f} and test loss: {best_loss:.4f}")
+    
+
+if __name__ == "__main__":
+    train1(NUM_SAMPLES=1000 ,maxNumShape=1, EPOCHS=50)
+    train1(NUM_SAMPLES=1000 ,maxNumShape=1, EPOCHS=50)
+    train1(NUM_SAMPLES=1000 ,maxNumShape=1, EPOCHS=50)
+    train1(NUM_SAMPLES=10000 ,maxNumShape=3, EPOCHS=50)
+    train1(NUM_SAMPLES=10000 ,maxNumShape=3, EPOCHS=50)
+    train1(NUM_SAMPLES=10000 ,maxNumShape=5, EPOCHS=50)
+    while True:
+        train1(NUM_SAMPLES=100000 ,maxNumShape=10, EPOCHS=5)
+    
+    

optimizer.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:820f2334048c67a59ebe18ff8b4b749a5a679d0e9facde598af5b697adee0f20
-size 34845043
+oid sha256:415436b37bc12580dca1104c9e1c4696dd84dd3926ff4af982670d2b2f8a0594
+size 144695175