Delete imageAI.py

imageAI.py

@@ -1,319 +0,0 @@
-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
-
-
-def generate_sample(num_shapes=1):
-    image = np.zeros((IMAGE_SIZE, IMAGE_SIZE), dtype=np.uint8)
-    instructions = []
-    
-    #num_shapes = random.randint(1, 3)
-    for _ in range(num_shapes):
-        shape = random.choice(['line', 'rectangle', 'circle', 'ellipse', 'polygon'])
-        color = random.randint(0, 255)
-        thickness = random.randint(1, 3)
-        
-        if shape == 'line':
-            start_point = (random.randint(0, IMAGE_SIZE), random.randint(0, IMAGE_SIZE))
-            end_point = (random.randint(0, IMAGE_SIZE), random.randint(0, IMAGE_SIZE))
-            cv2.line(image, start_point, end_point, color, thickness)
-            instructions.append(f"cv2.line(image, {start_point}, {end_point}, {color}, {thickness})")
-        
-        elif 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, thickness)
-            instructions.append(f"cv2.rectangle(image, {start_point}, {end_point}, {color}, {thickness})")
-        
-        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, thickness)
-            instructions.append(f"cv2.circle(image, {center}, {radius}, {color}, {thickness})")
-        
-        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, thickness)
-            instructions.append(f"cv2.ellipse(image, {center}, {axes}, {angle}, 0, 360, {color}, {thickness})")
-        
-        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.polylines(image, [points], True, color, thickness)
-            instructions.append(f"cv2.polylines(image, [{points.tolist()}], True, {color}, {thickness})")
-        
-    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']).unsqueeze(0) / 255.0
-        return image, len(sample['instructions'])
-
-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
-
-def train(model, train_loader, optimizer, criterion):
-    model.train()
-    total_loss = 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()
-        if batch_idx % 100 == 0:
-            print(f'Train Batch {batch_idx}/{len(train_loader)} Loss: {loss.item():.6f}')
-    return total_loss / len(train_loader)
-
-def test(model, test_loader, criterion, print_predictions=False):
-    model.eval()
-    test_loss = 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()
-            all_predictions.extend(output.cpu().numpy())
-            all_targets.extend(target.cpu().numpy())
-    
-    test_loss /= len(test_loader)
-    print(f'Test set: Average loss: {test_loss:.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, 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')
-
-    for epoch in range(EPOCHS):
-        print(f'Epoch {epoch+1}/{EPOCHS}')
-        train_loss = train(model, train_loader, optimizer, criterion)
-        test_loss, predictions, targets = test(model, test_loader, criterion, print_predictions=True)
-        
-        if test_loss < best_loss:
-            best_loss = test_loss
-            torch.save(model.state_dict(), os.path.join(Gbase, 'best_model.pth'))
-            torch.save(optimizer.state_dict(), os.path.join(Gbase, 'optimizer.pth'))
-            print(f"New best model saved with test loss: {best_loss:.4f}")
-
-
-
-
-def main():
-    # Set random seed
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = SimpleModel(path=Gbase+ 'best_model.pth').to(device)
-    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
-    if os.path.exists(Gbase+'optimizer.pth'):
-        print("Loading optimizer state...")
-        optimizer.load_state_dict(torch.load('optimizer.pth'))
-    criterion = nn.MSELoss()
-    test_image =Gbase+"image.jpg"
-    # np.random.randint(0, 256, (IMAGE_SIZE, IMAGE_SIZE), dtype=np.uint8)
-    instructions = model.predict(test_image)
-    print("Generated instructions:")
-    for instruction in instructions:
-        print(instruction)
-    # 檢查是否存在已保存的優化器狀態
-   
-    #return 
-    seed = 618 * 382 * 33
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-
-    # Generate dataset
-    dataset = generate_dataset()
-
-    # Split dataset into train and test
-    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')
-
-    for epoch in range(EPOCHS):
-        print(f'Epoch {epoch+1}/{EPOCHS}')
-        train_loss = train(model, train_loader, optimizer, criterion, device)
-        test_loss, predictions, targets = test(model, test_loader, criterion, device, print_predictions=True)
-        
-        if test_loss < best_loss:
-            best_loss = test_loss
-            torch.save(model.state_dict(),Gbase+ 'best_model.pth')
-            torch.save(optimizer.state_dict(),Gbase+ 'optimizer.pth')
-            print(f"New best model saved with test loss: {best_loss:.4f}")
-
-    # 測試 predict 方法
-    
-
-if __name__ == "__main__":
-    train1(NUM_SAMPLES=1000 ,maxNumShape=1, EPOCHS=100)
-    train1(NUM_SAMPLES=1000 ,maxNumShape=1, EPOCHS=100)
-    train1(NUM_SAMPLES=1000 ,maxNumShape=1, EPOCHS=100)
-    train1(NUM_SAMPLES=10000 ,maxNumShape=2, EPOCHS=10)
-    train1(NUM_SAMPLES=10000 ,maxNumShape=3,  EPOCHS=10)
-    train1(NUM_SAMPLES=100000 ,maxNumShape=5, EPOCHS=10)
-    train1(NUM_SAMPLES=100000 ,maxNumShape=5, EPOCHS=10)
-    train1(NUM_SAMPLES=100000 ,maxNumShape=5, EPOCHS=10)
-    while True:
-        train1(NUM_SAMPLES=100000 ,maxNumShape=8, EPOCHS=10)
-    
-