shancx 1.8.92__py3-none-any.whl → 1.9.33.218__py3-none-any.whl

shancx/Train/prepare.py

@@ -0,0 +1,161 @@
+import argparse
+from io import BytesIO
+import multiprocessing
+from multiprocessing import Lock, Process, RawValue
+from functools import partial
+from multiprocessing.sharedctypes import RawValue
+from PIL import Image
+from tqdm import tqdm
+from torchvision.transforms import functional as trans_fn
+import os
+from pathlib import Path
+import lmdb
+import numpy as np
+import time
+def resize_and_convert(img, size, resample):
+    if(img.size[0] != size):
+        img = trans_fn.resize(img, size, resample)
+        img = trans_fn.center_crop(img, size)
+    return img
+def image_convert_bytes(img):
+    buffer = BytesIO()
+    img.save(buffer, format='png')
+    return buffer.getvalue()
+def resize_multiple(img, sizes=(16, 128), resample=Image.BICUBIC, lmdb_save=False):
+    lr_img = resize_and_convert(img, sizes[0], resample)
+    hr_img = resize_and_convert(img, sizes[1], resample)
+    sr_img = resize_and_convert(lr_img, sizes[1], resample)
+    if lmdb_save:
+        lr_img = image_convert_bytes(lr_img)
+        hr_img = image_convert_bytes(hr_img)
+        sr_img = image_convert_bytes(sr_img)
+    return [lr_img, hr_img, sr_img]
+def resize_worker(img_file, sizes, resample, lmdb_save=False):
+    img = Image.open(img_file)
+    img = img.convert('RGB')
+    out = resize_multiple(
+        img, sizes=sizes, resample=resample, lmdb_save=lmdb_save)
+    return img_file.name.split('.')[0], out
+class WorkingContext():
+    def __init__(self, resize_fn, lmdb_save, out_path, env, sizes):
+        self.resize_fn = resize_fn
+        self.lmdb_save = lmdb_save
+        self.out_path = out_path
+        self.env = env
+        self.sizes = sizes
+        self.counter = RawValue('i', 0)
+        self.counter_lock = Lock()
+    def inc_get(self):
+        with self.counter_lock:
+            self.counter.value += 1
+            return self.counter.value
+    def value(self):
+        with self.counter_lock:
+            return self.counter.value
+def prepare_process_worker(wctx, file_subset):
+    for file in file_subset:
+        i, imgs = wctx.resize_fn(file)
+        lr_img, hr_img, sr_img = imgs
+        if not wctx.lmdb_save:
+            lr_img.save(
+                '{}/lr_{}/{}.png'.format(wctx.out_path, wctx.sizes[0], i.zfill(5)))
+            hr_img.save(
+                '{}/hr_{}/{}.png'.format(wctx.out_path, wctx.sizes[1], i.zfill(5)))
+            sr_img.save(
+                '{}/sr_{}_{}/{}.png'.format(wctx.out_path, wctx.sizes[0], wctx.sizes[1], i.zfill(5)))
+        else:
+            with wctx.env.begin(write=True) as txn:
+                txn.put('lr_{}_{}'.format(
+                    wctx.sizes[0], i.zfill(5)).encode('utf-8'), lr_img)
+                txn.put('hr_{}_{}'.format(
+                    wctx.sizes[1], i.zfill(5)).encode('utf-8'), hr_img)
+                txn.put('sr_{}_{}_{}'.format(
+                    wctx.sizes[0], wctx.sizes[1], i.zfill(5)).encode('utf-8'), sr_img)
+        curr_total = wctx.inc_get()
+        if wctx.lmdb_save:
+            with wctx.env.begin(write=True) as txn:
+                txn.put('length'.encode('utf-8'), str(curr_total).encode('utf-8'))
+def all_threads_inactive(worker_threads):
+    for thread in worker_threads:
+        if thread.is_alive():
+            return False
+    return True
+def prepare(img_path, out_path, n_worker=3, sizes=(16, 128), resample=Image.BICUBIC, lmdb_save=False):
+    resize_fn = partial(resize_worker, sizes=sizes,
+                        resample=resample, lmdb_save=lmdb_save)
+    files = [p for p in Path(
+        '{}'.format(img_path)).glob(f'**/*')]
+
+    if not lmdb_save:
+        os.makedirs(out_path, exist_ok=True)
+        os.makedirs('{}/lr_{}'.format(out_path, sizes[0]), exist_ok=True)
+        os.makedirs('{}/hr_{}'.format(out_path, sizes[1]), exist_ok=True)
+        os.makedirs('{}/sr_{}_{}'.format(out_path,
+                    sizes[0], sizes[1]), exist_ok=True)
+    else:
+        env = lmdb.open(out_path, map_size=1024 ** 4, readahead=False)
+    if n_worker > 1:
+        # prepare data subsets
+        multi_env = None
+        if lmdb_save:
+            multi_env = env
+        file_subsets = np.array_split(files, n_worker)
+        worker_threads = []
+        wctx = WorkingContext(resize_fn, lmdb_save, out_path, multi_env, sizes)
+        # start worker processes, monitor results
+        for i in range(n_worker):
+            proc = Process(target=prepare_process_worker, args=(wctx, file_subsets[i]))
+            proc.start()
+            worker_threads.append(proc)        
+        total_count = str(len(files))
+        while not all_threads_inactive(worker_threads):
+            print("\r{}/{} images processed".format(wctx.value(), total_count), end=" ")
+            time.sleep(0.1)
+    else:
+        total = 0
+        for file in tqdm(files):
+            i, imgs = resize_fn(file)
+            lr_img, hr_img, sr_img = imgs
+            if not lmdb_save:
+                lr_img.save(
+                    '{}/lr_{}/{}.png'.format(out_path, sizes[0], i.zfill(5)))
+                hr_img.save(
+                    '{}/hr_{}/{}.png'.format(out_path, sizes[1], i.zfill(5)))
+                sr_img.save(
+                    '{}/sr_{}_{}/{}.png'.format(out_path, sizes[0], sizes[1], i.zfill(5)))
+            else:
+                with env.begin(write=True) as txn:
+                    txn.put('lr_{}_{}'.format(
+                        sizes[0], i.zfill(5)).encode('utf-8'), lr_img)
+                    txn.put('hr_{}_{}'.format(
+                        sizes[1], i.zfill(5)).encode('utf-8'), hr_img)
+                    txn.put('sr_{}_{}_{}'.format(
+                        sizes[0], sizes[1], i.zfill(5)).encode('utf-8'), sr_img)
+            total += 1
+            if lmdb_save:
+                with env.begin(write=True) as txn:
+                    txn.put('length'.encode('utf-8'), str(total).encode('utf-8'))
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--path', '-p', type=str,
+                        default='{}/Dataset/celebahq_256'.format(Path.home()))
+    parser.add_argument('--out', '-o', type=str,
+                        default='./dataset/celebahq')
+    parser.add_argument('--size', type=str, default='64,512')
+    parser.add_argument('--n_worker', type=int, default=3)
+    parser.add_argument('--resample', type=str, default='bicubic')
+    # default save in png format
+    parser.add_argument('--lmdb', '-l', action='store_true')
+
+    args = parser.parse_args()
+
+    resample_map = {'bilinear': Image.BILINEAR, 'bicubic': Image.BICUBIC}
+    resample = resample_map[args.resample]
+    sizes = [int(s.strip()) for s in args.size.split(',')]
+
+    args.out = '{}_{}_{}'.format(args.out, sizes[0], sizes[1])
+
+    #1.输入 输出文件夹
+    prepare(args.path, args.out, args.n_worker,
+            sizes=sizes, resample=resample, lmdb_save=args.lmdb)

shancx/Train/renet50.py

@@ -0,0 +1,157 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+import torch
+import torchvision
+import torchvision.transforms as transforms
+import cv2
+import numpy as np
+from typing import List, Tuple
+from PIL import Image
+import os
+
+# class ModelLoader:
+#     """Handles loading and managing the pretrained model"""
+    
+#     def __init__(self, model_name: str = 'resnet50'):
+#         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+#         self.model = self._load_pretrained_model(model_name)
+        
+#     def _load_pretrained_model(self, model_name: str) -> torch.nn.Module:
+#         """Load pretrained model from torchvision.models"""
+#         model = getattr(torchvision.models, model_name)(pretrained=True)
+#         return model.eval().to(self.device)
+
+class ModelLoader:
+    """Handles loading and managing the pretrained model"""
+
+    def __init__(self, model_name: str = 'resnet50'):
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model_dir = "models" 
+        os.makedirs(self.model_dir, exist_ok=True) 
+        self.model = self._load_pretrained_model(model_name)
+
+    def _load_pretrained_model(self, model_name: str) -> torch.nn.Module:
+        """Load pretrained model from torchvision.models. If not found locally, download and save it."""
+        model_path = os.path.join(self.model_dir, f"{model_name}.pth")  
+
+        try:
+            model = getattr(torchvision.models, model_name)(pretrained=False)  
+            model.load_state_dict(torch.load(model_path))  
+            print(f"Loaded {model_name} from local directory: {model_path}")
+        except FileNotFoundError:
+            print(f"Model {model_name} not found locally. Downloading...")
+            model = getattr(torchvision.models, model_name)(pretrained=True)  
+            torch.save(model.state_dict(), model_path)  
+            print(f"Downloaded and saved {model_name} to: {model_path}")
+        except Exception as e:
+            raise RuntimeError(f"Failed to load or download the model: {e}")
+
+        return model.eval().to(self.device)
+    
+
+class ImageProcessor:
+    """Handles image preprocessing and transformations"""
+    
+    def __init__(self):
+        self.transform = transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225]
+            )
+        ])
+        
+    def preprocess_image(self, image_path: str) -> torch.Tensor:
+        """Load and preprocess image for model input"""
+        try:
+            if not os.path.exists(image_path):
+                raise FileNotFoundError(f"Image file not found at {image_path}")
+                
+            image = cv2.imread(image_path)
+            if image is None:
+                raise ValueError(f"Unable to read image at {image_path}")
+                
+            # Convert BGR to RGB
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            # Convert to PIL Image for torchvision transforms
+            image = Image.fromarray(image)
+            return self.transform(image).unsqueeze(0)
+        except Exception as e:
+            print(f"Error processing image: {e}")
+            raise
+
+class Predictor:
+    """Handles model predictions and visualization"""
+    
+    def __init__(self, model_loader: ModelLoader, labels_path: str = 'imagenet_classes.txt'):
+        self.model = model_loader.model
+        self.device = model_loader.device
+        self.labels = self._load_labels(labels_path)
+        
+    def _load_labels(self, labels_path: str) -> List[str]:
+        """Load class labels from file"""
+        try:
+            with open(labels_path) as f:
+                labels = [line.strip() for line in f.readlines()]
+                if len(labels) != 1000:
+                    raise ValueError(f"Expected 1000 ImageNet classes, got {len(labels)}")
+                return labels
+        except Exception as e:
+            print(f"Error loading labels: {e}")
+            raise
+            
+    def predict(self, input_tensor: torch.Tensor) -> Tuple[int, str]:
+        """Run model prediction on input tensor"""
+        with torch.no_grad():
+            input_tensor = input_tensor.to(self.device)
+            pred = self.model(input_tensor)
+            pred_index = torch.argmax(pred, 1).cpu().detach().numpy()[0]
+            
+            if pred_index >= len(self.labels):
+                raise ValueError(f"Prediction index {pred_index} out of range for {len(self.labels)} classes")
+                
+            return pred_index, self.labels[pred_index]
+            
+    def visualize_prediction(self, image_path: str, class_name: str):
+        """Display image with predicted class label"""
+        try:
+            image = cv2.imread(image_path)
+            if image is None:
+                raise FileNotFoundError(f"Image not found at {image_path}")
+                
+            cv2.putText(image, class_name, (50, 50), 
+                       cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2)
+            # cv2.imshow("Prediction", image)
+            # cv2.waitKey(0)
+            # cv2.destroyAllWindows()
+            output_path = f"output_image_{class_name}.jpg"  # 替换为你想保存的路径
+            cv2.imwrite(output_path, image)
+            print(f"Image saved to {output_path}")
+        except Exception as e:
+            print(f"Error visualizing prediction: {e}")
+            raise
+
+def main():
+    try:
+        # Initialize components
+        model_loader = ModelLoader()
+        image_processor = ImageProcessor()
+        predictor = Predictor(model_loader)
+        
+        # Process image and make prediction
+        image_path = "./space_shuttle.jpg"
+        input_tensor = image_processor.preprocess_image(image_path)
+        pred_index, class_name = predictor.predict(input_tensor)
+        
+        # Display results
+        print(f"Predicted class index: {pred_index}")
+        print(f"Predicted class name: {class_name}")
+        predictor.visualize_prediction(image_path, class_name)
+        
+    except Exception as e:
+        print(f"Error in main execution: {e}")
+
+# if __name__ == "__main__":
+#     main()

shancx/ZR.py

@@ -0,0 +1,12 @@
+def dbz2rfl(d):  # 全局定义
+    return 10. ** (d / 10.)
+
+def rfl2mmh(z, a=200., b=1.6):  # 全局定义
+    return (z / a) ** (1. / b)
+
+def ZR1(conf):
+    dbz = conf[0]
+    z_values = dbz2rfl(dbz)  # 直接调用全局函数
+    rainfall_mmh = rfl2mmh(z_values)
+    rainfall_mmh[rainfall_mmh < 0.1] = 0
+    return rainfall_mmh