diff --git a/demo/demo/seg_demo.cpp b/demo/demo/seg_demo.cpp
index 658e0b9..35b4fa8 100644
--- a/demo/demo/seg_demo.cpp
+++ b/demo/demo/seg_demo.cpp
@@ -14,20 +14,18 @@ void sig_handler(int signo) {
     gRun = false;
 }
 
-void writePred(const std::string& images_names, const std::string& gt_folder, const std::string& out_folder, tk::dnn::SegmentationNN& segNN){
+void writePred(const std::string& images_names, const std::string& gt_folder, const std::string& out_folder, tk::dnn::SegmentationNN& segNN, int& width, int& height, bool show=false){
     std::ifstream all_gt(images_names);
     std::string filename;
     cv::Mat frame;
-    std::vector<cv::Mat> batch_frame;
-    std::vector<cv::Mat> batch_dnn_input;
     for (; std::getline(all_gt, filename); ) {
         std::cout<<filename<<std::endl;
         frame = cv::imread(gt_folder + filename);
-        batch_dnn_input.clear();
-        batch_frame.clear();
-        batch_frame.push_back(frame);
-        batch_dnn_input.push_back(frame.clone());
-        segNN.update(batch_dnn_input, 1, false);
+        height = frame.rows;
+        width = frame.cols;
+        segNN.updateOriginal(frame, false);
+        if(show)
+            segNN.draw();
         cv::imwrite(out_folder + filename, segNN.segmented[0]);
     }
 }
@@ -50,7 +48,7 @@ int main(int argc, char *argv[]) {
     int n_classes = 19;
     if(argc > 4)
         n_classes = atoi(argv[4]); 
-    bool show = false;
+    bool show = true;
     if(argc > 5)
         show = atoi(argv[5]); 
     bool write_pred = false;
@@ -64,80 +62,65 @@ int main(int argc, char *argv[]) {
     tk::dnn::SegmentationNN segNN;
     segNN.init(net, n_classes, n_batch);
 
+    int height = 0, width = 0;
+
     if(write_pred){
         std::string gt_folder = "../demo/CityScapes_val/images/";
         std::string images_names = "../demo/CityScapes_val/all_images.txt";
         std::string out_folder = "seg/";
 
-        writePred(images_names, gt_folder, out_folder, segNN);
-        return 0;
+        writePred(images_names, gt_folder, out_folder, segNN, width, height, show);
     }
+    else{
+        if(!show)
+            SAVE_RESULT = true;
 
-    if(!show)
-        SAVE_RESULT = true;
+        gRun = true;
 
-    gRun = true;
+        cv::VideoCapture cap(input);
+        if(!cap.isOpened())
+            gRun = false; 
+        else
+            std::cout<<"camera started\n";
 
-    cv::VideoCapture cap(input);
-    if(!cap.isOpened())
-        gRun = false; 
-    else
-        std::cout<<"camera started\n";
+        cv::VideoWriter resultVideo;
+        if(SAVE_RESULT) {
+            int w = cap.get(cv::CAP_PROP_FRAME_WIDTH);
+            int h = cap.get(cv::CAP_PROP_FRAME_HEIGHT);
+            resultVideo.open("result.mp4", cv::VideoWriter::fourcc('M','P','4','V'), 30, cv::Size(1024, 1024));
+        }
 
-    cv::VideoWriter resultVideo;
-    if(SAVE_RESULT) {
-        int w = cap.get(cv::CAP_PROP_FRAME_WIDTH);
-        int h = cap.get(cv::CAP_PROP_FRAME_HEIGHT);
-        resultVideo.open("result.mp4", cv::VideoWriter::fourcc('M','P','4','V'), 30, cv::Size(1024, 1024));
-    }
-
-    cv::Mat frame;
-    std::vector<cv::Mat> batch_frame;
-    std::vector<cv::Mat> batch_dnn_input;
-    int height = 0, width = 0;
-
-    while(gRun) {
-        batch_dnn_input.clear();
-        batch_frame.clear();
-        
-        for(int bi=0; bi< n_batch; ++bi){
+        cv::Mat frame;
+        while(gRun) {
             cap >> frame; 
             if(!frame.data) 
                 break;
             height = frame.rows;
             width = frame.cols;
-            batch_frame.push_back(frame);
 
-            // this will be resized to the net format
-            batch_dnn_input.push_back(frame.clone());
-        } 
-        if(!frame.data) 
-            break;
-    
-        //inference
-        segNN.update(batch_dnn_input, n_batch);
-        frame = segNN.draw();
+            //inference
+            segNN.updateOriginal(frame);
+            if(show)
+                segNN.draw();
 
-        if(n_batch == 1 && SAVE_RESULT)
-            resultVideo << frame;
+            if(SAVE_RESULT)
+                resultVideo << segNN.segmented[0];
+        }
     }
 
     std::cout<<"segmentation end\n";   
     double mean = 0, mean_pre = 0, mean_post = 0;
     
     std::cout<<COL_GREENB<<"\n\nTime stats for size ["<<width<<","<<height<<"] :\n";
-    // std::cout<<"Min: "<<*std::min_element(segNN.stats.begin(), segNN.stats.end())/n_batch<<" ms\n";    
-    // std::cout<<"Max: "<<*std::max_element(segNN.stats.begin(), segNN.stats.end())/n_batch<<" ms\n";    
     
     for(int i=0; i<segNN.stats.size(); i++) mean += segNN.stats[i]; mean /= segNN.stats.size();
     for(int i=0; i<segNN.stats_pre.size(); i++) mean_pre += segNN.stats_pre[i]; mean_pre /= segNN.stats_pre.size();
     for(int i=0; i<segNN.stats_post.size(); i++) mean_post += segNN.stats_post[i]; mean_post /= segNN.stats_post.size();
-    std::cout<<"Avg pre:\t"<<mean_pre/n_batch<<" ms\t"<<1000/(mean_pre/n_batch)<<" FPS\n";   
-    std::cout<<"Avg inf:\t"<<mean/n_batch<<" ms\t"<<1000/(mean/n_batch)<<" FPS\n";   
-    std::cout<<"Avg post:\t"<<mean_post/n_batch<<" ms\t"<<1000/(mean_post/n_batch)<<" FPS\n\n";   
-    std::cout<<"Avg tot:\t"<<(mean_pre + mean_post + mean) /n_batch<<" ms\t"<<1000/((mean_pre + mean_post + mean)/n_batch)<<" FPS\n"<<COL_END;   
+    std::cout<<"Avg pre:\t"<<mean_pre<<" ms\t"<<1000/(mean_pre)<<" FPS\n";   
+    std::cout<<"Avg inf:\t"<<mean<<" ms\t"<<1000/(mean)<<" FPS\n";   
+    std::cout<<"Avg post:\t"<<mean_post<<" ms\t"<<1000/(mean_post)<<" FPS\n\n";   
+    std::cout<<"Avg tot:\t"<<(mean_pre + mean_post + mean) <<" ms\t"<<1000/((mean_pre + mean_post + mean))<<" FPS\n"<<COL_END;   
     
-
     return 0;
 }
 
diff --git a/include/tkDNN/SegmentationNN.h b/include/tkDNN/SegmentationNN.h
index e4f5729..edede43 100644
--- a/include/tkDNN/SegmentationNN.h
+++ b/include/tkDNN/SegmentationNN.h
@@ -232,7 +232,156 @@ class SegmentationNN {
                 TKDNN_TSTOP
                 stats_post.push_back(t_ns);
             }
-        }      
+        }     
+
+        void updateOriginal(cv::Mat frame, bool apply_colormap=true){
+
+            std::vector<cv::Mat> splitted_frames;
+            int H, W, net_H, net_W;
+            int top = 0, bottom = 0, left = 0, right = 0;
+            std::vector<std::pair<int,int>> pos;
+
+            {
+                TKDNN_TSTART
+                cv::Size original_size =  frame.size();
+
+                frame.convertTo(frame, CV_32FC3, 1 / 255.0, 0);
+                H = frame.rows;
+                W = frame.cols;
+                net_H = netRT->input_dim.h; 
+                net_W = netRT->input_dim.w;
+
+                cv::Mat frame_cropped;
+                
+                if( H <= net_H && W <= net_W ){ // smaller size wrt network
+                    top = (net_H - H)/2;
+                    bottom = net_H - H - top ;
+                    left = (net_W - W)/2;
+                    right = net_W - W - left ;
+                    cv::copyMakeBorder(frame, frame_cropped, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(0,0,0) );
+                    splitted_frames.push_back(frame_cropped);
+                }
+                else{ //bigger size wrt network
+
+                    
+                    if(H < net_H  ||  W < net_W){
+                        if(H < net_H){
+                            top = (net_H - H)/2;
+                            bottom = net_H - H - top ;
+                        }
+                        else{
+                            left = (net_W - W)/2;
+                            right = net_W - W - left ;
+                        }
+                        cv::copyMakeBorder(frame, frame_cropped, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(0,0,0));
+                    }
+
+                    for(int x=0; x+net_W<=W ;){
+                        for(int y=0; y+net_H <=H ; ){
+                            cv::Rect roi(x, y, net_W, net_H);
+                            cv::Mat image_roi = frame(roi);
+                            splitted_frames.push_back(image_roi);
+                            pos.push_back(std::make_pair(x,y));
+                            
+                            y += net_H;
+                            if(y == H)
+                                break;
+                            if(y + net_H > H) y = H - net_H;
+                        }
+                        x += net_W;
+                        if(x == W)
+                            break;
+                        if(x + net_W > W) x = W - net_W;
+                    }
+                }
+
+                tk::dnn::dataDim_t idim = netRT->input_dim;
+
+                if(splitted_frames.size()> nBatches)
+                    FatalError(std::to_string(splitted_frames.size()) + " min batches required");
+
+                for(int bi=0; bi<splitted_frames.size();++bi){
+                    cv::split(splitted_frames[bi], bgr);
+                    for (int i = 0; i < idim.c; i++){
+                        int idx = i * splitted_frames[bi].rows * splitted_frames[bi].cols;
+                        int ch = idim.c-1 -i;
+                        memcpy((void *)&input[idx + idim.tot()*bi], (void *)bgr[ch].data, splitted_frames[bi].rows * splitted_frames[bi].cols * sizeof(dnnType));
+                    }
+
+                    checkCuda(cudaMemcpyAsync(input_d+ idim.tot()*bi, input + idim.tot()*bi, idim.tot() * sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
+                    normalize(input_d + idim.tot()*bi, idim.c, idim.h, idim.w, mean_d, stddev_d);
+                }
+                TKDNN_TSTOP
+                stats_pre.push_back(t_ns);
+            }
+
+            tk::dnn::dataDim_t dim = netRT->input_dim;
+            dim.n = splitted_frames.size();
+            {
+                if(TKDNN_VERBOSE) dim.print();
+                TKDNN_TSTART
+                netRT->infer(dim, input_d);
+                TKDNN_TSTOP
+                if(TKDNN_VERBOSE) dim.print();
+                stats.push_back(t_ns);
+            }
+
+            dataDim_t odim = netRT->output_dim;
+
+            std::vector<cv::Mat> out_img;
+
+            {
+                TKDNN_TSTART
+
+                for(int bi=0; bi<splitted_frames.size();++bi){
+
+                    dnnType *rt_out = (dnnType *)netRT->buffersRT[1]+ netRT->buffersDIM[1].tot()*bi;               
+                    
+                    matrixTranspose(cublasHandle, rt_out, tmpInputData_d, odim.c, odim.w*odim.h);
+                    maxElem(tmpInputData_d, tmpOutData_d, odim.c, odim.h, odim.w);
+                    checkCuda(cudaMemcpy(tmpOutData_h, tmpOutData_d, odim.w*odim.h * sizeof(float), cudaMemcpyDeviceToHost));
+
+                    dataDim_t vdim = odim;
+                    vdim.c = 1;
+
+                    cv::Mat colored;
+
+                    if(apply_colormap)
+                        colored = vizData2Mat(tmpOutData_h, vdim, 1024, 0, 18);
+                    else{
+                        cv::Mat colored_fp32 (cv::Size(odim.w, odim.h),CV_32FC1, tmpOutData_h);
+                        colored_fp32.convertTo(colored, CV_8UC1);
+                    }
+                    out_img.push_back(colored);
+                }
+
+
+                cv::Mat seg(frame.size(), out_img[0].type());
+                if(out_img.size() == 1)
+                {
+                    cv::Rect roi(left, top, W, H);
+                    seg = out_img[0](roi);
+                }
+                else{
+                    int bi=0;
+
+                    if(top == 0 && left == 0){
+
+                        for(int i=0; i<out_img.size(); ++i){
+                            cv::Mat roi_collage = seg(cv::Rect( pos[i].first ,pos[i].second,out_img[i].cols,out_img[i].rows));
+                            out_img[i].copyTo(roi_collage);
+                        }
+                    }
+                    else{
+                        FatalError("Not handled case")
+                    }
+                }
+                segmented[0] = seg;
+                
+                TKDNN_TSTOP
+                stats_post.push_back(t_ns);
+            }
+        } 
 
         /**
          * Method to draw boundixg boxes and labels on a frame.
@@ -240,9 +389,9 @@ class SegmentationNN {
         cv::Mat draw(const int cur_batches=1) {
             for(int i=0; i<cur_batches; ++i){
 
-                // cv::imshow("segmented", segmented[i]);
-                // cv::resizeWindow("segmented", cv::Size(512,288));
-                // cv::waitKey(1);
+                cv::imshow("segmented", segmented[i]);
+                cv::resizeWindow("segmented", cv::Size(512,288));
+                cv::waitKey(1);
             }
             return segmented[0];
         }