Create imageAI.py

imageAI.py

@@ -0,0 +1,108 @@
+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)
+except:
+  IN_COLAB = False
+  Gbase="./"
+  cache_dir="./hf/"
+
+
+import cv2,os
+import numpy as np
+import random,string
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+IMAGE_SIZE = 64
+NUM_SAMPLES = 1000
+BATCH_SIZE = 4
+EPOCHS = 500
+LEARNING_RATE = 0.001
+
+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.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.dropout(x)
+        x = F.leaky_relu(self.fc2(x))
+        x = self.dropout(x)
+        x = self.fc3(x)
+        return x
+
+    def predict(self, image):
+        self.eval()
+        with torch.no_grad():
+            if isinstance(image, str) and os.path.isfile(image):
+                # 如果輸入是圖片文件路徑
+                img = cv2.imread(image, cv2.IMREAD_GRAYSCALE)
+                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)
+                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 = 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)})")
+                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)})")
+                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)})")
+                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)})")
+              
+            
+            return instructions
+